| 166 | |
| 167 | |
| 168 | class EncoderBlockRes1B(nn.Module): |
| 169 | def __init__( |
| 170 | self, |
| 171 | in_channels: int, |
| 172 | out_channels: int, |
| 173 | kernel_size: Tuple, |
| 174 | downsample: Tuple, |
| 175 | momentum: float, |
| 176 | has_film, |
| 177 | ): |
| 178 | r"""Encoder block, contains 8 convolutional layers.""" |
| 179 | super(EncoderBlockRes1B, self).__init__() |
| 180 | |
| 181 | self.conv_block1 = ConvBlockRes( |
| 182 | in_channels, out_channels, kernel_size, momentum, has_film, |
| 183 | ) |
| 184 | self.downsample = downsample |
| 185 | |
| 186 | def forward(self, input_tensor: torch.Tensor, film_dict: Dict) -> torch.Tensor: |
| 187 | r"""Forward data into the module. |
| 188 | |
| 189 | Args: |
| 190 | input_tensor: (batch_size, input_feature_maps, time_steps, freq_bins) |
| 191 | |
| 192 | Returns: |
| 193 | encoder_pool: (batch_size, output_feature_maps, downsampled_time_steps, downsampled_freq_bins) |
| 194 | encoder: (batch_size, output_feature_maps, time_steps, freq_bins) |
| 195 | """ |
| 196 | encoder = self.conv_block1(input_tensor, film_dict['conv_block1']) |
| 197 | encoder_pool = F.avg_pool2d(encoder, kernel_size=self.downsample) |
| 198 | return encoder_pool, encoder |
| 199 | |
| 200 | |
| 201 | class DecoderBlockRes1B(nn.Module): |