| 142 | |
| 143 | |
| 144 | class AudioEncoder(nn.Module): |
| 145 | def __init__(self, n_frames): |
| 146 | super().__init__() |
| 147 | self.n_frames = n_frames |
| 148 | self.first_net = nn.Sequential( |
| 149 | ConvNormRelu(1, 64, '2d', False), |
| 150 | ConvNormRelu(64, 64, '2d', True), |
| 151 | ConvNormRelu(64, 128, '2d', False), |
| 152 | ConvNormRelu(128, 128, '2d', True), |
| 153 | ConvNormRelu(128, 256, '2d', False), |
| 154 | ConvNormRelu(256, 256, '2d', True), |
| 155 | ConvNormRelu(256, 256, '2d', False), |
| 156 | ConvNormRelu(256, 256, '2d', False, padding='VALID') |
| 157 | ) |
| 158 | |
| 159 | self.make_1d = torch.nn.Upsample((n_frames, 1), mode='bilinear', align_corners=False) |
| 160 | |
| 161 | self.down1 = nn.Sequential( |
| 162 | ConvNormRelu(256, 256, '1d', False), |
| 163 | ConvNormRelu(256, 256, '1d', False) |
| 164 | ) |
| 165 | self.down2 = ConvNormRelu(256, 256, '1d', True) |
| 166 | self.down3 = ConvNormRelu(256, 256, '1d', True) |
| 167 | self.down4 = ConvNormRelu(256, 256, '1d', True) |
| 168 | self.down5 = ConvNormRelu(256, 256, '1d', True) |
| 169 | self.down6 = ConvNormRelu(256, 256, '1d', True) |
| 170 | self.up1 = UnetUp(256, 256) |
| 171 | self.up2 = UnetUp(256, 256) |
| 172 | self.up3 = UnetUp(256, 256) |
| 173 | self.up4 = UnetUp(256, 256) |
| 174 | self.up5 = UnetUp(256, 256) |
| 175 | |
| 176 | def forward(self, spectrogram): |
| 177 | spectrogram = spectrogram.unsqueeze(1) # add channel dim |
| 178 | # print(spectrogram.shape) |
| 179 | spectrogram = spectrogram.float() |
| 180 | |
| 181 | out = self.first_net(spectrogram) |
| 182 | out = self.make_1d(out) |
| 183 | x1 = out.squeeze(3) |
| 184 | |
| 185 | x2 = self.down1(x1) |
| 186 | x3 = self.down2(x2) |
| 187 | x4 = self.down3(x3) |
| 188 | x5 = self.down4(x4) |
| 189 | x6 = self.down5(x5) |
| 190 | x7 = self.down6(x6) |
| 191 | x = self.up1(x7, x6) |
| 192 | x = self.up2(x, x5) |
| 193 | x = self.up3(x, x4) |
| 194 | x = self.up4(x, x3) |
| 195 | x = self.up5(x, x2) |
| 196 | |
| 197 | return x |
| 198 | |
| 199 | |
| 200 | class Generator(nn.Module): |