I’ve got the following mechanism for adding positional awareness to representations, it works well but is very slow.
The fundamental issue is that each sample in the batch has a different real length and I therefore have to loop over and create different embeddings for each of them.
class PositionSymmetricLayer(nn.Module):
def __init__(self, dim):
super().__init__()
self.embedding = nn.Parameter(torch.linspace(0., 1., dim//2), requires_grad=False)
self.dim = dim
def forward(self, x, durs):
N = x.size(0)
T = x.size(1)
pos_emb = torch.zeros((N, T, self.dim), dtype=torch.float32).cuda()
for i, dur in enumerate(durs):
factor = torch.linspace(0, 1, dur).unsqueeze(1).cuda()
pos = (torch.cat((self.embedding - factor, -self.embedding + factor,), dim=-1) >= 0).float()
pos_emb[i, -dur:] = pos
x = torch.cat((x, pos_emb,), dim=-1)
return x
To help understand what this does you can see in the following what pos
is given a dur
of 3 and a dim
of 8:
tensor([[1., 1., 1., 1., 1., 0., 0., 0.],
[0., 0., 1., 1., 1., 1., 0., 0.],
[0., 0., 0., 1., 1., 1., 1., 1.]])
Is there something I can do to speed the for loop up?