以下是使用PyTorch实现Transformer的代码：
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
def __init__(self, heads, d_model):
super().__init__()
self.heads = heads
self.d_model = d_model
self.d_head = d_model // heads
self.q_linear = nn.Linear(d_model, d_model)
self.v_linear = nn.Linear(d_model, d_model)
self.k_linear = nn.Linear(d_model, d_model)
self.out = nn.Linear(d_model, d_model)
def forward(self, q, k, v, mask=None):
bs = q.size(0)
# Linear projections
k = self.k_linear(k).view(bs, -1, self.heads, self.d_head)
q = self.q_linear(q).view(bs, -1, self.heads, self.d_head)
v = self.v_linear(v).view(bs, -1, self.heads, self.d_head)
# Transpose and dot product attention
k = k.transpose(1,2)
q = q.transpose(1,2)
v = v.transpose(1,2)
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head)
if mask is not None:
mask = mask.unsqueeze(1)
scores = scores.masked_fill(mask == 0, -1e9)
scores = F.softmax(scores, dim=-1)
# Output attention
output = torch.matmul(scores, v)
# Concatenate and linear projection
output = output.transpose(1,2).contiguous().view(bs, -1, self.d_model)
return self.out(output)
class PositionwiseFeedforward(nn.Module):
def __init__(self, d_model, d_ff=2048):
super().__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.linear2 = nn.Linear(d_ff, d_model)
def forward(self, x):
x = self.linear1(x)
x = F.relu(x)
x = self.linear2(x)
return x
class EncoderLayer(nn.Module):
def __init__(self, d_model, heads, dropout=0.1):
super().__init__()
self.norm_1 = nn.LayerNorm(d_model)
self.norm_2 = nn.LayerNorm(d_model)
self.attn = MultiHeadAttention(heads, d_model)
self.ff = PositionwiseFeedforward(d_model)
self.dropout_1 = nn.Dropout(dropout)
self.dropout_2 = nn.Dropout(dropout)
def forward(self, x, mask):
x2 = self.norm_1(x)
x = x + self.dropout_1(self.attn(x2, x2, x2, mask))
x2 = self.norm_2(x)
x = x + self.dropout_2(self.ff(x2))
return x
class TransformerEncoder(nn.Module):
def __init__(self, input_dim, d_model, heads, num_layers):
super().__init__()
self.input_dim = input_dim
self.d_model = d_model
self.heads = heads
self.num_layers = num_layers
self.embedding = nn.Embedding(input_dim, d_model)
self.pe = PositionalEncoder(d_model)
self.layers = nn.ModuleList([EncoderLayer(d_model, heads) for _ in range(num_layers)])
def forward(self, src_seq, src_mask):
x = self.embedding(src_seq)
x = self.pe(x)
for i in range(self.num_layers):
x = self.layers[i](x, src_mask)
return x
class PositionalEncoder(nn.Module):
def __init__(self, d_model, dropout=0.1, max_len=5000):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0).transpose(0, 1)
self.register_buffer('pe', pe)
def forward(self, x):
x = x + self.pe[:x.size(0), :]
return self.dropout(x)
class Transformer(nn.Module):
def __init__(self, input_dim, output_dim, d_model, heads, num_layers, dropout=0.1):
super().__init__()
self.encoder = TransformerEncoder(input_dim, d_model, heads, num_layers)
self.fc = nn.Linear(d_model, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, src_seq, src_mask):
x = self.encoder(src_seq, src_mask)
x = x.mean(dim=1)
x = self.fc(x)
return x
```
以上是使用PyTorch实现Transformer的代码，其中包括了Transformer的编码器、多头自注意力机制、位置编码等模块。需要注意的是，该代码中使用了Layer Normalization进行层归一化。
1482
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
