Upsampling block.
| 514 | return x |
| 515 | |
| 516 | class PSUpsamplerBlock(nn.Module): |
| 517 | """ |
| 518 | Upsampling block. |
| 519 | """ |
| 520 | def __init__(self, d_model, d_model_out, scale_factor, resolution=None, view_num=1): |
| 521 | super().__init__() |
| 522 | |
| 523 | self.scale_factor = scale_factor |
| 524 | self.d_model_out = d_model_out |
| 525 | self.residual_fc = nn.Linear(d_model, d_model_out * (scale_factor**2), bias=True) |
| 526 | self.pixelshuffle = nn.PixelShuffle(scale_factor) |
| 527 | |
| 528 | self.resolution = resolution |
| 529 | |
| 530 | self.view_num = view_num |
| 531 | |
| 532 | def forward(self, x): |
| 533 | x = self.residual_fc(x) |
| 534 | bs, l, c = x.shape |
| 535 | x = x.reshape(bs * self.view_num, self.resolution[0], self.resolution[1], c).permute(0, 3, 1, 2) |
| 536 | x = self.pixelshuffle(x) |
| 537 | x = x.permute(0, 2, 3, 1).reshape(bs, self.view_num, self.resolution[0]*self.scale_factor, self.resolution[1]*self.scale_factor, self.d_model_out) |
| 538 | x = x.reshape(bs, -1, self.d_model_out) |
| 539 | return x |
| 540 | |
| 541 | class GaussianUpsampler(nn.Module): |
| 542 | """ |