37 lines
962 B
Python
37 lines
962 B
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
__author__ = "Yu-Hsiang Huang"
|
|
|
|
|
|
class ScaledDotProductAttention(nn.Module):
|
|
"""Scaled Dot-Product Attention"""
|
|
|
|
def __init__(self, temperature, attn_dropout=0.1):
|
|
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 == 0, -1e9)
|
|
|
|
attn = self.dropout(F.softmax(attn, dim=-1))
|
|
output = torch.matmul(attn, v)
|
|
|
|
return output, attn
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def forward(self, query, key, value):
|
|
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query.size(-1))
|
|
|
|
p_attn = F.softmax(scores, dim=-1)
|
|
p_val = torch.matmul(p_attn, value)
|
|
return p_val, p_attn
|