(self, x, feat_cache=None, feat_idx=[0])
| 562 | ) |
| 563 | |
| 564 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 565 | |
| 566 | if feat_cache is not None: |
| 567 | idx = feat_idx[0] |
| 568 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 569 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 570 | cache_x = torch.cat( |
| 571 | [ |
| 572 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 573 | cache_x.device), |
| 574 | cache_x, |
| 575 | ], |
| 576 | dim=2, |
| 577 | ) |
| 578 | x = self.conv1(x, feat_cache[idx]) |
| 579 | feat_cache[idx] = cache_x |
| 580 | feat_idx[0] += 1 |
| 581 | else: |
| 582 | x = self.conv1(x) |
| 583 | |
| 584 | ## downsamples |
| 585 | for layer in self.downsamples: |
| 586 | if feat_cache is not None: |
| 587 | x = layer(x, feat_cache, feat_idx) |
| 588 | else: |
| 589 | x = layer(x) |
| 590 | |
| 591 | ## middle |
| 592 | for layer in self.middle: |
| 593 | if isinstance(layer, ResidualBlock) and feat_cache is not None: |
| 594 | x = layer(x, feat_cache, feat_idx) |
| 595 | else: |
| 596 | x = layer(x) |
| 597 | |
| 598 | ## head |
| 599 | for layer in self.head: |
| 600 | if isinstance(layer, CausalConv3d) and feat_cache is not None: |
| 601 | idx = feat_idx[0] |
| 602 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 603 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 604 | cache_x = torch.cat( |
| 605 | [ |
| 606 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 607 | cache_x.device), |
| 608 | cache_x, |
| 609 | ], |
| 610 | dim=2, |
| 611 | ) |
| 612 | x = layer(x, feat_cache[idx]) |
| 613 | feat_cache[idx] = cache_x |
| 614 | feat_idx[0] += 1 |
| 615 | else: |
| 616 | x = layer(x) |
| 617 | |
| 618 | return x |
| 619 | |
| 620 | |
| 621 | class Decoder3d(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected