Upsampler.
| 539 | return x |
| 540 | |
| 541 | class GaussianUpsampler(nn.Module): |
| 542 | """ |
| 543 | Upsampler. |
| 544 | """ |
| 545 | def __init__(self, width, up_ratio, ch_decay=1, low_channels=64, window_size=None, opt: Options=None): |
| 546 | super().__init__() |
| 547 | self.up_ratio = up_ratio |
| 548 | self.low_channels = low_channels |
| 549 | self.window_size = window_size |
| 550 | |
| 551 | self.base_width = width |
| 552 | |
| 553 | if len(opt.down_resolution) > 0: |
| 554 | self.input_res = (opt.down_resolution[0] // opt.patch_size, opt.down_resolution[1] // opt.patch_size) |
| 555 | else: |
| 556 | self.input_res = (opt.image_height // opt.patch_size, opt.image_width // opt.patch_size) |
| 557 | |
| 558 | resolution = [self.input_res[0], self.input_res[1]] |
| 559 | for res_log2 in range(int(np.log2(up_ratio))): |
| 560 | _width = width |
| 561 | width = max(width // ch_decay, self.low_channels) |
| 562 | heads = int(width / 64) |
| 563 | width = heads * 64 |
| 564 | self.add_module(f'upsampler_{res_log2}', PSUpsamplerBlock(_width, width, scale_factor=2, resolution=resolution, view_num=opt.input_frames)) |
| 565 | resolution = [resolution[0]*2, resolution[1]*2] |
| 566 | encoder = Transformer(width, 2, heads, window_size=window_size) |
| 567 | self.add_module(f'attention_{res_log2}', encoder) |
| 568 | self.out_channels = width |
| 569 | self.layernorm2 = LayerNorm(width) |
| 570 | |
| 571 | def forward(self, x): |
| 572 | for res_log2 in range(int(np.log2(self.up_ratio))): |
| 573 | x = getattr(self, f'upsampler_{res_log2}')(x) |
| 574 | x = getattr(self, f'attention_{res_log2}')(x) |
| 575 | x = self.layernorm2(x) |
| 576 | return x |
| 577 | |
| 578 | |
| 579 |