| 727 | return h |
| 728 | |
| 729 | class Upsample3D(nn.Module): |
| 730 | def __init__(self, |
| 731 | in_channels, |
| 732 | scale_factor=2 |
| 733 | ): |
| 734 | super().__init__() |
| 735 | |
| 736 | self.scale_factor = scale_factor |
| 737 | self.conv3d = Res3DBlockUpsample(input_filters=in_channels, |
| 738 | num_filters=in_channels, |
| 739 | down_sampling_stride=(1, 1, 1), |
| 740 | down_sampling=False) |
| 741 | |
| 742 | def forward(self, x, is_init=True, is_split=True): |
| 743 | b, c, t, h, w = x.shape |
| 744 | |
| 745 | # x = x.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3).to(memory_format=torch.channels_last_3d) |
| 746 | if is_split: |
| 747 | split_size = c // 8 |
| 748 | x_slices = torch.split(x, split_size, dim=1) |
| 749 | x = [nn.functional.interpolate(x, scale_factor=self.scale_factor) for x in x_slices] |
| 750 | x = torch.cat(x, dim=1) |
| 751 | else: |
| 752 | x = nn.functional.interpolate(x, scale_factor=self.scale_factor) |
| 753 | |
| 754 | x = self.conv3d(x, is_init) |
| 755 | return x |
| 756 | |
| 757 | class VideoDecoder(nn.Module): |
| 758 | def __init__(self, |