55 lines
1.2 KiB
Python
55 lines
1.2 KiB
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class PositionalEncoding(nn.Module):
|
|
"""
|
|
Standard sinusoidal positional encoding from
|
|
"Attention Is All You Need".
|
|
"""
|
|
|
|
def __init__(self, embed_dim, dropout=0.1, max_len=10000):
|
|
super().__init__()
|
|
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
pe = torch.zeros(max_len, embed_dim)
|
|
|
|
position = torch.arange(
|
|
0,
|
|
max_len,
|
|
dtype=torch.float
|
|
).unsqueeze(1)
|
|
|
|
div_term = torch.exp(
|
|
torch.arange(
|
|
0,
|
|
embed_dim,
|
|
2
|
|
).float() *
|
|
(-math.log(10000.0) / embed_dim)
|
|
)
|
|
|
|
pe[:, 0::2] = torch.sin(position * div_term)
|
|
pe[:, 1::2] = torch.cos(position * div_term)
|
|
|
|
pe = pe.unsqueeze(0)
|
|
|
|
# Stored as a buffer so it moves with the model
|
|
# but isn't trained.
|
|
self.register_buffer("pe", pe)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Args:
|
|
x: Tensor of shape
|
|
(batch_size, sequence_length, embed_dim)
|
|
"""
|
|
|
|
seq_len = x.size(1)
|
|
|
|
x = x + self.pe[:, :seq_len]
|
|
|
|
return self.dropout(x) |