70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
|
import math
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
|
||
|
|
|
||
|
|
class PositionalEncoding(nn.Module):
|
||
|
|
"""Sinusoidal positional encoding for non-recurrent neural networks.
|
||
|
|
Implementation based on "Attention Is All You Need"
|
||
|
|
|
||
|
|
Args:
|
||
|
|
channels (int): embedding size
|
||
|
|
dropout_p (float): dropout rate applied to the output.
|
||
|
|
max_len (int): maximum sequence length.
|
||
|
|
use_scale (bool): whether to use a learnable scaling coefficient.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, channels, dropout_p=0.0, max_len=5000, use_scale=False):
|
||
|
|
super().__init__()
|
||
|
|
if channels % 2 != 0:
|
||
|
|
raise ValueError(
|
||
|
|
"Cannot use sin/cos positional encoding with " "odd channels (got channels={:d})".format(channels)
|
||
|
|
)
|
||
|
|
self.use_scale = use_scale
|
||
|
|
if use_scale:
|
||
|
|
self.scale = torch.nn.Parameter(torch.ones(1))
|
||
|
|
pe = torch.zeros(max_len, channels)
|
||
|
|
position = torch.arange(0, max_len).unsqueeze(1)
|
||
|
|
div_term = torch.pow(10000, torch.arange(0, channels, 2).float() / channels)
|
||
|
|
pe[:, 0::2] = torch.sin(position.float() * div_term)
|
||
|
|
pe[:, 1::2] = torch.cos(position.float() * div_term)
|
||
|
|
pe = pe.unsqueeze(0).transpose(1, 2)
|
||
|
|
self.register_buffer("pe", pe)
|
||
|
|
if dropout_p > 0:
|
||
|
|
self.dropout = nn.Dropout(p=dropout_p)
|
||
|
|
self.channels = channels
|
||
|
|
|
||
|
|
def forward(self, x, mask=None, first_idx=None, last_idx=None):
|
||
|
|
"""
|
||
|
|
Shapes:
|
||
|
|
x: [B, C, T]
|
||
|
|
mask: [B, 1, T]
|
||
|
|
first_idx: int
|
||
|
|
last_idx: int
|
||
|
|
"""
|
||
|
|
|
||
|
|
x = x * math.sqrt(self.channels)
|
||
|
|
if first_idx is None:
|
||
|
|
if self.pe.size(2) < x.size(2):
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Sequence is {x.size(2)} but PositionalEncoding is"
|
||
|
|
f" limited to {self.pe.size(2)}. See max_len argument."
|
||
|
|
)
|
||
|
|
if mask is not None:
|
||
|
|
pos_enc = self.pe[:, :, : x.size(2)] * mask
|
||
|
|
else:
|
||
|
|
pos_enc = self.pe[:, :, : x.size(2)]
|
||
|
|
if self.use_scale:
|
||
|
|
x = x + self.scale * pos_enc
|
||
|
|
else:
|
||
|
|
x = x + pos_enc
|
||
|
|
else:
|
||
|
|
if self.use_scale:
|
||
|
|
x = x + self.scale * self.pe[:, :, first_idx:last_idx]
|
||
|
|
else:
|
||
|
|
x = x + self.pe[:, :, first_idx:last_idx]
|
||
|
|
if hasattr(self, "dropout"):
|
||
|
|
x = self.dropout(x)
|
||
|
|
return x
|