| 5 | |
| 6 | class HeatMapEncoder(nn.Module): |
| 7 | def __init__(self, dim_in=768, dim_aud=64, win_size=16): |
| 8 | super(HeatMapEncoder, self).__init__() |
| 9 | self.win_size = win_size |
| 10 | self.dim_aud = dim_aud |
| 11 | |
| 12 | num_layers = 5 |
| 13 | kernel_size = 3 |
| 14 | dilation_rate = 2 |
| 15 | |
| 16 | lst = [] |
| 17 | for i in range(num_layers): |
| 18 | dilation = dilation_rate ** i |
| 19 | padding = int((kernel_size * dilation - dilation) / 2) |
| 20 | if i == 0: |
| 21 | in_dim = dim_in |
| 22 | out_dim = 16 |
| 23 | else: |
| 24 | in_dim = 16 |
| 25 | out_dim = 16 |
| 26 | lst.append(nn.Conv2d(in_dim, out_dim, kernel_size, dilation=dilation, padding=padding)) |
| 27 | lst.append(nn.LeakyReLU(0.02, True)) |
| 28 | self.conv = nn.Sequential(*lst) |
| 29 | |
| 30 | self.fc = nn.Sequential( |
| 31 | nn.Linear(dim_in, 32), |
| 32 | # nn.Linear(16, 32), |
| 33 | nn.LeakyReLU(0.02, True), |
| 34 | nn.Linear(32, 32), |
| 35 | nn.LeakyReLU(0.02, True), |
| 36 | nn.Linear(32, 64), |
| 37 | nn.LeakyReLU(0.02, True), |
| 38 | nn.Linear(64, 32), |
| 39 | nn.LeakyReLU(0.02, True), |
| 40 | # nn.Linear(64, 64), |
| 41 | # nn.LeakyReLU(0.02, True), |
| 42 | nn.Linear(32, dim_aud), |
| 43 | ) |
| 44 | def forward(self, x): |
| 45 | """ |
| 46 | x: [b, 1, n_rays, c] |