在PyTorch中實現Transformer模型需要定義Transformer的各個組件,包括Encoder、Decoder、Multihead Attention、Feedforward等。以下是一個簡單的Transformer模型的實現示例:
import torch
import torch.nn as nn
import torch.nn.functional as F
# 定義Multihead Attention層
class MultiheadAttention(nn.Module):
def __init__(self, d_model, n_head):
super(MultiheadAttention, self).__init__()
self.d_model = d_model
self.n_head = n_head
self.head_dim = d_model // n_head
self.fc_q = nn.Linear(d_model, d_model)
self.fc_k = nn.Linear(d_model, d_model)
self.fc_v = nn.Linear(d_model, d_model)
self.fc_o = nn.Linear(d_model, d_model)
def forward(self, q, k, v):
q = self.fc_q(q)
k = self.fc_k(k)
v = self.fc_v(v)
q = q.view(q.size(0), -1, self.n_head, self.head_dim).transpose(1, 2)
k = k.view(k.size(0), -1, self.n_head, self.head_dim).transpose(1, 2)
v = v.view(v.size(0), -1, self.n_head, self.head_dim).transpose(1, 2)
attention = F.softmax(torch.matmul(q, k.transpose(-2, -1)) / self.head_dim, dim=-1)
output = torch.matmul(attention, v).transpose(1, 2).contiguous().view(q.size(0), -1, self.d_model)
output = self.fc_o(output)
return output
# 定義Feedforward層
class Feedforward(nn.Module):
def __init__(self, d_model, d_ff):
super(Feedforward, self).__init__()
self.fc1 = nn.Linear(d_model, d_ff)
self.fc2 = nn.Linear(d_ff, d_model)
def forward(self, x):
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
# 定義Encoder層
class EncoderLayer(nn.Module):
def __init__(self, d_model, n_head, d_ff):
super(EncoderLayer, self).__init__()
self.multihead_attention = MultiheadAttention(d_model, n_head)
self.feedforward = Feedforward(d_model, d_ff)
def forward(self, x):
att_output = self.multihead_attention(x, x, x)
ff_output = self.feedforward(att_output)
output = x + att_output + ff_output
return output
# 定義Transformer模型
class Transformer(nn.Module):
def __init__(self, d_model, n_head, d_ff, num_layers):
super(Transformer, self).__init__()
self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, n_head, d_ff) for _ in range(num_layers)])
def forward(self, x):
for encoder_layer in self.encoder_layers:
x = encoder_layer(x)
return x
# 使用Transformer模型
d_model = 512
n_head = 8
d_ff = 2048
num_layers = 6
transformer = Transformer(d_model, n_head, d_ff, num_layers)
input_data = torch.randn(10, 20, d_model)
output = transformer(input_data)
print(output.size())
在這個示例中,我們定義了Multihead Attention層、Feedforward層、EncoderLayer和Transformer模型,并使用這些組件來構建一個簡單的Transformer模型。您可以根據具體的任務和需求對模型進行調整和修改。