| 212 | |
| 213 | |
| 214 | class AttModule(nn.Module): |
| 215 | def __init__(self, dilation, in_channels, out_channels, r1, r2, att_type, stage, alpha): |
| 216 | super(AttModule, self).__init__() |
| 217 | self.feed_forward = ConvFeedForward(dilation, in_channels, out_channels) |
| 218 | self.instance_norm = nn.InstanceNorm1d(in_channels, track_running_stats=False) |
| 219 | self.att_layer = AttLayer(in_channels, in_channels, out_channels, r1, r1, r2, dilation, att_type=att_type, stage=stage) # dilation |
| 220 | self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1) |
| 221 | self.dropout = nn.Dropout() |
| 222 | self.alpha = alpha |
| 223 | |
| 224 | def forward(self, x, f, mask): |
| 225 | out = self.feed_forward(x) |
| 226 | out = self.alpha * self.att_layer(self.instance_norm(out), f, mask) + out |
| 227 | out = self.conv_1x1(out) |
| 228 | out = self.dropout(out) |
| 229 | return (x + out) * mask[:, 0:1, :] |
| 230 | |
| 231 | |
| 232 | class PositionalEncoding(nn.Module): |