我已经实现了https://arxiv.org/pdf/1710.10903.pdf的注意力(方程式 1),但它显然不是内存效率,并且只能在我的 GPU 上运行一个模型(需要 7-10GB)。
目前,我有
class MyModule(nn.Module):
def __init__(self, in_features, out_features):
super(MyModule, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.W = nn.Parameter(nn.init.xavier_uniform(torch.Tensor(in_features, out_features).type(torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor), gain=np.sqrt(2.0)), requires_grad=True)
self.a = nn.Parameter(nn.init.xavier_uniform(torch.Tensor(2*out_features, 1).type(torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor), gain=np.sqrt(2.0)), requires_grad=True)
def forward(self, input):
h = torch.mm(input, self.W)
N = h.size()[0]
a_input = torch.cat([h.repeat(1, N).view(N * N, -1), h.repeat(N, 1)], dim=1).view(N, -1, 2 * self.out_features)
e = F.elu(torch.matmul(a_input, self.a).squeeze(2))
return e
我计算所有 e_ij 项的见解是
In [8]: import torch
在 [9] 中:将 numpy 导入为 np
在 [10] 中: h = torch.LongTensor(np.array([[1,1], [2,2], [3,3]]))
在[11]中:N=3
在 [12] 中:h.repeat(1, N).view(N * N, -1) 出[12]:
1 1
1 1
1 1
2 2
2 2
2 2
3 3
3 3
3 3
[torch.LongTensor 大小为 9x2]
在 [13] 中:h.repeat(N, 1) 出 [13]:
1 1
2 2
3 3
1 1
2 2
3 3
1 1
2 2
3 3
[torch.LongTensor 大小为 9x2]
最后连接 hs 和 feed 矩阵 a。
有没有办法以更记忆友好的方式做到这一点?