| 158 | |
| 159 | |
| 160 | class Upsample2D(nn.Module): |
| 161 | def __init__(self, |
| 162 | channels, |
| 163 | use_conv=False, |
| 164 | use_conv_transpose=False, |
| 165 | out_channels=None): |
| 166 | super().__init__() |
| 167 | self.channels = channels |
| 168 | self.out_channels = out_channels or channels |
| 169 | self.use_conv = use_conv |
| 170 | self.use_conv_transpose = use_conv_transpose |
| 171 | |
| 172 | if use_conv: |
| 173 | self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1) |
| 174 | else: |
| 175 | assert "Not Supported" |
| 176 | self.conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1) |
| 177 | |
| 178 | def forward(self, x, output_size=None): |
| 179 | assert x.shape[-1] == self.channels |
| 180 | |
| 181 | if self.use_conv_transpose: |
| 182 | return self.conv(x) |
| 183 | |
| 184 | if output_size is None: |
| 185 | x = F.interpolate( |
| 186 | x.permute(0,3,1,2).to(memory_format=torch.channels_last), |
| 187 | scale_factor=2.0, mode='nearest').permute(0,2,3,1).contiguous() |
| 188 | else: |
| 189 | x = F.interpolate( |
| 190 | x.permute(0,3,1,2).to(memory_format=torch.channels_last), |
| 191 | size=output_size, mode='nearest').permute(0,2,3,1).contiguous() |
| 192 | |
| 193 | # x = self.conv(x) |
| 194 | x = base_conv2d(x, self.conv, channel_last=True) |
| 195 | return x |
| 196 | |
| 197 | |
| 198 | class Downsample2D(nn.Module): |