| 52 | return x * y |
| 53 | |
| 54 | class audioEncoder(nn.Module): |
| 55 | def __init__(self, layers, num_filters, **kwargs): |
| 56 | super(audioEncoder, self).__init__() |
| 57 | block = SEBasicBlock |
| 58 | self.inplanes = num_filters[0] |
| 59 | |
| 60 | self.conv1 = nn.Conv2d(1, num_filters[0] , kernel_size=7, stride=(2, 1), padding=3, |
| 61 | bias=False) |
| 62 | self.bn1 = nn.BatchNorm2d(num_filters[0]) |
| 63 | self.relu = nn.ReLU(inplace=True) |
| 64 | |
| 65 | self.layer1 = self._make_layer(block, num_filters[0], layers[0]) |
| 66 | self.layer2 = self._make_layer(block, num_filters[1], layers[1], stride=(2, 2)) |
| 67 | self.layer3 = self._make_layer(block, num_filters[2], layers[2], stride=(2, 2)) |
| 68 | self.layer4 = self._make_layer(block, num_filters[3], layers[3], stride=(1, 1)) |
| 69 | out_dim = num_filters[3] * block.expansion |
| 70 | |
| 71 | for m in self.modules(): |
| 72 | if isinstance(m, nn.Conv2d): |
| 73 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') |
| 74 | elif isinstance(m, nn.BatchNorm2d): |
| 75 | nn.init.constant_(m.weight, 1) |
| 76 | nn.init.constant_(m.bias, 0) |
| 77 | |
| 78 | def _make_layer(self, block, planes, blocks, stride=1): |
| 79 | downsample = None |
| 80 | if stride != 1 or self.inplanes != planes * block.expansion: |
| 81 | downsample = nn.Sequential( |
| 82 | nn.Conv2d(self.inplanes, planes * block.expansion, |
| 83 | kernel_size=1, stride=stride, bias=False), |
| 84 | nn.BatchNorm2d(planes * block.expansion), |
| 85 | ) |
| 86 | |
| 87 | layers = [] |
| 88 | layers.append(block(self.inplanes, planes, stride, downsample)) |
| 89 | self.inplanes = planes * block.expansion |
| 90 | for i in range(1, blocks): |
| 91 | layers.append(block(self.inplanes, planes)) |
| 92 | |
| 93 | return nn.Sequential(*layers) |
| 94 | |
| 95 | def forward(self, x): |
| 96 | x = self.conv1(x) |
| 97 | x = self.bn1(x) |
| 98 | x = self.relu(x) |
| 99 | |
| 100 | x = self.layer1(x) |
| 101 | x = self.layer2(x) |
| 102 | x = self.layer3(x) |
| 103 | x = self.layer4(x) |
| 104 | x = torch.mean(x, dim=2, keepdim=True) |
| 105 | x = x.view((x.size()[0], x.size()[1], -1)) |
| 106 | x = x.transpose(1, 2) |
| 107 | |
| 108 | return x |