I’m trying to implement the attention module with 1D convolution.
Here’s my custom attention module for that.
class convSelfAttention(nn.Module):
def __init__(self, in_dim):
super(convSelfAttention, self).__init__()
self.channel_in = in_dim
self.query_conv = nn.Conv1d(in_channels=in_dim, out_channels=in_dim//2, kernel_size=1)
self.key_conv = nn.Conv1d(in_channels=in_dim, out_channels=in_dim//2, kernel_size=1)
self.value_conv = nn.Conv1d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
self.gamma = nn.Parameter(torch.zeros(1))
self.activ = nn.Softmax(dim=-1)
def forward(self, x):
p_query = self.query_conv(x).permute(0, 2, 1)
p_key = self.key_conv(x)
energy = torch.bmm(p_query, p_key)
attention = self.activ(energy)
p_value = self.value_conv(x)
out = torch.bmm(p_value, attention.permute(0, 2, 1))
out = self.gamma * out + x
out = out.permute(0, 2, 1)
return out
However, the gradients are 0 except ‘self.gamma’.
For your convenience, I plot the gradient flow of my entire architecture.(The last two layers are fc and all layers from ‘~~Encoder’ are convolutions)
