(self, x)
| 333 | np.concatenate((H1r, G1r), axis=0))) |
| 334 | |
| 335 | def forward(self, x): |
| 336 | B, N, c, k = x.shape # (B, N, k) |
| 337 | ns = math.floor(np.log2(N)) |
| 338 | nl = pow(2, math.ceil(np.log2(N))) |
| 339 | extra_x = x[:, 0:nl - N, :, :] |
| 340 | x = torch.cat([x, extra_x], 1) |
| 341 | Ud = torch.jit.annotate(List[Tensor], []) |
| 342 | Us = torch.jit.annotate(List[Tensor], []) |
| 343 | # decompose |
| 344 | for i in range(ns - self.L): |
| 345 | # print('x shape',x.shape) |
| 346 | d, x = self.wavelet_transform(x) |
| 347 | Ud += [self.A(d) + self.B(x)] |
| 348 | Us += [self.C(d)] |
| 349 | x = self.T0(x) # coarsest scale transform |
| 350 | |
| 351 | # reconstruct |
| 352 | for i in range(ns - 1 - self.L, -1, -1): |
| 353 | x = x + Us[i] |
| 354 | x = torch.cat((x, Ud[i]), -1) |
| 355 | x = self.evenOdd(x) |
| 356 | x = x[:, :N, :, :] |
| 357 | |
| 358 | return x |
| 359 | |
| 360 | def wavelet_transform(self, x): |
| 361 | xa = torch.cat([x[:, ::2, :, :], |
nothing calls this directly
no test coverage detected