File size: 1,270 Bytes
cb0ad2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
import math
import torch
from torch import nn
class SALayer(nn.Module):
def __init__(self, in_dim, att_dim, head_nums):
super().__init__()
self.in_dim = in_dim
self.att_dim = att_dim
self.head_nums = head_nums
assert self.in_dim % self.head_nums == 0
self.key_layer = nn.Conv1d(self.in_dim, self.att_dim * self.head_nums, 1, 1, 0)
self.query_layer = nn.Conv1d(self.in_dim, self.att_dim * self.head_nums, 1, 1, 0)
self.value_layer = nn.Conv1d(self.in_dim, self.in_dim, 1, 1, 0)
self.scale = 1 / math.sqrt(self.att_dim)
def forward(self, feats, masks=None):
bs, c, n = feats.shape
keys = self.key_layer(feats).reshape(bs, -1, self.head_nums, n)
querys = self.query_layer(feats).reshape(bs, -1, self.head_nums, n)
values = self.value_layer(feats).reshape(bs, -1, self.head_nums, n)
logits = torch.einsum('bchk,bchq->bhkq', keys, querys) * self.scale
if masks is not None:
logits = logits - (1 - masks[:, None, :, None]) * 1e8
weights = torch.softmax(logits, dim=2)
new_feats = torch.einsum('bchk,bhkq->bchq', values, weights)
new_feats = new_feats.reshape(bs, -1, n)
return new_feats + feats
|