| 199 | |
| 200 | |
| 201 | class DecoderBlockRes1B(nn.Module): |
| 202 | def __init__( |
| 203 | self, |
| 204 | in_channels: int, |
| 205 | out_channels: int, |
| 206 | kernel_size: Tuple, |
| 207 | upsample: Tuple, |
| 208 | momentum: float, |
| 209 | has_film, |
| 210 | ): |
| 211 | r"""Decoder block, contains 1 transposed convolutional and 8 convolutional layers.""" |
| 212 | super(DecoderBlockRes1B, self).__init__() |
| 213 | self.kernel_size = kernel_size |
| 214 | self.stride = upsample |
| 215 | |
| 216 | self.conv1 = torch.nn.ConvTranspose2d( |
| 217 | in_channels=in_channels, |
| 218 | out_channels=out_channels, |
| 219 | kernel_size=self.stride, |
| 220 | stride=self.stride, |
| 221 | padding=(0, 0), |
| 222 | bias=False, |
| 223 | dilation=(1, 1), |
| 224 | ) |
| 225 | |
| 226 | self.bn1 = nn.BatchNorm2d(in_channels, momentum=momentum) |
| 227 | self.conv_block2 = ConvBlockRes( |
| 228 | out_channels * 2, out_channels, kernel_size, momentum, has_film, |
| 229 | ) |
| 230 | self.bn2 = nn.BatchNorm2d(in_channels, momentum=momentum) |
| 231 | self.has_film = has_film |
| 232 | |
| 233 | self.init_weights() |
| 234 | |
| 235 | def init_weights(self): |
| 236 | r"""Initialize weights.""" |
| 237 | init_bn(self.bn1) |
| 238 | init_layer(self.conv1) |
| 239 | |
| 240 | def forward( |
| 241 | self, input_tensor: torch.Tensor, concat_tensor: torch.Tensor, film_dict: Dict, |
| 242 | ) -> torch.Tensor: |
| 243 | r"""Forward data into the module. |
| 244 | |
| 245 | Args: |
| 246 | input_tensor: (batch_size, input_feature_maps, downsampled_time_steps, downsampled_freq_bins) |
| 247 | concat_tensor: (batch_size, input_feature_maps, time_steps, freq_bins) |
| 248 | |
| 249 | Returns: |
| 250 | output_tensor: (batch_size, output_feature_maps, time_steps, freq_bins) |
| 251 | """ |
| 252 | # b1 = film_dict['beta1'] |
| 253 | |
| 254 | b1 = film_dict['beta1'] |
| 255 | x = self.conv1(F.leaky_relu_(self.bn1(input_tensor) + b1)) |
| 256 | # (batch_size, input_feature_maps, time_steps, freq_bins) |
| 257 | |
| 258 | x = torch.cat((x, concat_tensor), dim=1) |