| 173 | |
| 174 | |
| 175 | class SPConvTranspose2d(nn.Module): |
| 176 | def __init__(self, in_channels, out_channels, kernel_size, r=1): |
| 177 | # upconvolution only along second dimension of image |
| 178 | # Upsampling using sub pixel layers |
| 179 | super(SPConvTranspose2d, self).__init__() |
| 180 | self.out_channels = out_channels |
| 181 | self.conv = nn.Conv2d(in_channels, out_channels * r, kernel_size=kernel_size, stride=(1, 1)) |
| 182 | self.r = r |
| 183 | |
| 184 | def forward(self, x): |
| 185 | out = self.conv(x) |
| 186 | batch_size, nchannels, H, W = out.shape |
| 187 | out = out.view((batch_size, self.r, nchannels // self.r, H, W)) |
| 188 | out = out.permute(0, 2, 3, 4, 1) |
| 189 | out = out.contiguous().view((batch_size, nchannels // self.r, H, -1)) |
| 190 | return out |
| 191 | |
| 192 | |
| 193 | class DenseBlock(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected