(self, x, feat_cache=None, feat_idx=[0])
| 103 | self.resample = nn.Identity() |
| 104 | |
| 105 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 106 | b, c, t, h, w = x.size() |
| 107 | if self.mode == 'upsample3d': |
| 108 | if feat_cache is not None: |
| 109 | idx = feat_idx[0] |
| 110 | if feat_cache[idx] is None: |
| 111 | feat_cache[idx] = 'Rep' |
| 112 | feat_idx[0] += 1 |
| 113 | else: |
| 114 | |
| 115 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 116 | if cache_x.shape[2] < 2 and feat_cache[ |
| 117 | idx] is not None and feat_cache[idx] != 'Rep': |
| 118 | # cache last frame of last two chunk |
| 119 | cache_x = torch.cat([ |
| 120 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 121 | cache_x.device), cache_x |
| 122 | ], |
| 123 | dim=2) |
| 124 | if cache_x.shape[2] < 2 and feat_cache[ |
| 125 | idx] is not None and feat_cache[idx] == 'Rep': |
| 126 | cache_x = torch.cat([ |
| 127 | torch.zeros_like(cache_x).to(cache_x.device), |
| 128 | cache_x |
| 129 | ], |
| 130 | dim=2) |
| 131 | if feat_cache[idx] == 'Rep': |
| 132 | x = self.time_conv(x) |
| 133 | else: |
| 134 | x = self.time_conv(x, feat_cache[idx]) |
| 135 | feat_cache[idx] = cache_x |
| 136 | feat_idx[0] += 1 |
| 137 | |
| 138 | x = x.reshape(b, 2, c, t, h, w) |
| 139 | x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), |
| 140 | 3) |
| 141 | x = x.reshape(b, c, t * 2, h, w) |
| 142 | t = x.shape[2] |
| 143 | x = rearrange(x, 'b c t h w -> (b t) c h w') |
| 144 | x = self.resample(x) |
| 145 | x = rearrange(x, '(b t) c h w -> b c t h w', t=t) |
| 146 | |
| 147 | if self.mode == 'downsample3d': |
| 148 | if feat_cache is not None: |
| 149 | idx = feat_idx[0] |
| 150 | if feat_cache[idx] is None: |
| 151 | feat_cache[idx] = x.clone() |
| 152 | feat_idx[0] += 1 |
| 153 | else: |
| 154 | |
| 155 | cache_x = x[:, :, -1:, :, :].clone() |
| 156 | # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep': |
| 157 | # # cache last frame of last two chunk |
| 158 | # cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) |
| 159 | |
| 160 | x = self.time_conv( |
| 161 | torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) |
| 162 | feat_cache[idx] = cache_x |
nothing calls this directly
no outgoing calls
no test coverage detected