mirror of
https://github.com/wassname/pytorch-transformer-ts.git
synced 2026-08-11 11:24:32 +08:00
26 lines
639 B
Python
26 lines
639 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class ScaledDotProductAttention(nn.Module):
|
|
""" Scaled Dot-Product Attention """
|
|
|
|
def __init__(self, temperature, attn_dropout=0.2):
|
|
super().__init__()
|
|
|
|
self.temperature = temperature
|
|
self.dropout = nn.Dropout(attn_dropout)
|
|
|
|
def forward(self, q, k, v, mask=None):
|
|
attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
|
|
|
|
if mask is not None:
|
|
attn = attn.masked_fill(mask, -1e9)
|
|
|
|
attn = self.dropout(F.softmax(attn, dim=-1))
|
|
output = torch.matmul(attn, v)
|
|
|
|
return output, attn
|
|
|