- def forward(self, x):
- q = torch.einsum('ntc,hdc->nhtd', x, self.w_q)
- k = torch.einsum('ntc,hdc->nhtd', x, self.w_k)
- v = torch.einsum('ntc,hdc->nhtd', x, self.w_v)
- r = math.sqrt(q.size(3))
- a = torch.einsum('nhtd,nhsd->nhts', q, k).div(r)
+ def forward(self, x_q, x_kv = None):
+ if x_kv is None: x_kv = x_q
+ q = torch.einsum('ntc,hdc->nhtd', x_q, self.w_q)
+ k = torch.einsum('ntc,hdc->nhtd', x_kv, self.w_k)
+ v = torch.einsum('ntc,hdc->nhtd', x_kv, self.w_v)
+ a = torch.einsum('nhtd,nhsd->nhts', q, k) / math.sqrt(q.size(3))