Transformer的python代码
时间: 2023-07-23 08:03:32 浏览: 119
以下是一个简单的Transformer模型的Python代码示例:
```python
import torch
import torch.nn as nn
import torch.optim as optim
class Transformer(nn.Module):
def __init__(self, input_dim, output_dim, hidden_dim, num_layers, num_heads):
super(Transformer, self).__init__()
self.embedding = nn.Embedding(input_dim, hidden_dim)
self.positional_encoding = PositionalEncoding(hidden_dim)
self.encoder = Encoder(hidden_dim, num_layers, num_heads)
self.decoder = Decoder(hidden_dim, num_layers, num_heads)
self.fc = nn.Linear(hidden_dim, output_dim)
def forward(self, src, trg):
src_embedding = self.embedding(src)
src_embedding = self.positional_encoding(src_embedding)
trg_embedding = self.embedding(trg)
trg_embedding = self.positional_encoding(trg_embedding)
encoder_output = self.encoder(src_embedding)
decoder_output = self.decoder(trg_embedding, encoder_output)
output = self.fc(decoder_output)
return output
class PositionalEncoding(nn.Module):
def __init__(self, hidden_dim, max_len=1000):
super(PositionalEncoding, self).__init__()
position = torch.arange(0, max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, hidden_dim, 2) * (-math.log(10000.0) / hidden_dim))
pe = torch.zeros(max_len, hidden_dim)
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 x
class Encoder(nn.Module):
def __init__(self, hidden_dim, num_layers, num_heads):
super(Encoder, self).__init__()
self.layers = nn.ModuleList([EncoderLayer(hidden_dim, num_heads) for _ in range(num_layers)])
def forward(self, x):
for layer in self.layers:
x = layer(x)
return x
class EncoderLayer(nn.Module):
def __init__(self, hidden_dim, num_heads):
super(EncoderLayer, self).__init__()
self.multihead_attention = MultiheadAttention(hidden_dim, num_heads)
self.feed_forward = FeedForward(hidden_dim)
self.layer_norm = nn.LayerNorm(hidden_dim)
def forward(self, x):
attention_output = self.multihead_attention(x)
x = x + attention_output
x = self.layer_norm(x)
feed_forward_output = self.feed_forward(x)
x = x + feed_forward_output
x = self.layer_norm(x)
return x
class Decoder(nn.Module):
def __init__(self, hidden_dim, num_layers, num_heads):
super(Decoder, self).__init__()
self.layers = nn.ModuleList([DecoderLayer(hidden_dim, num_heads) for _ in range(num_layers)])
def forward(self, x, encoder_output):
for layer in self.layers:
x = layer(x, encoder_output)
return x
class DecoderLayer(nn.Module):
def __init__(self, hidden_dim, num_heads):
super(DecoderLayer, self).__init__()
self.masked_multihead_attention = MultiheadAttention(hidden_dim, num_heads)
self.multihead_attention = MultiheadAttention(hidden_dim, num_heads)
self.feed_forward = FeedForward(hidden_dim)
self.layer_norm = nn.LayerNorm(hidden_dim)
def forward(self, x, encoder_output):
masked_attention_output = self.masked_multihead_attention(x, x, x)
x = x + masked_attention_output
x = self.layer_norm(x)
attention_output = self.multihead_attention(x, encoder_output, encoder_output)
x = x + attention_output
x = self.layer_norm(x)
feed_forward_output = self.feed_forward(x)
x = x + feed_forward_output
x = self.layer_norm(x)
return x
class MultiheadAttention(nn.Module):
def __init__(self, hidden_dim, num_heads):
super(MultiheadAttention, self).__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads
self.query_linear = nn.Linear(hidden_dim, hidden_dim)
self.key_linear = nn.Linear(hidden_dim, hidden_dim)
self.value_linear = nn.Linear(hidden_dim, hidden_dim)
self.fc = nn.Linear(hidden_dim, hidden_dim)
def forward(self, query, key, value):
batch_size = query.shape[0]
query = self.query_linear(query)
key = self.key_linear(key)
value = self.value_linear(value)
query = query.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
key = key.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
value = value.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
attention_scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim)
attention_weights = nn.functional.softmax(attention_scores, dim=-1)
output = torch.matmul(attention_weights, value)
output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.hidden_dim)
output = self.fc(output)
return output
class FeedForward(nn.Module):
def __init__(self, hidden_dim, feed_forward_dim=2048):
super(FeedForward, self).__init__()
self.fc1 = nn.Linear(hidden_dim, feed_forward_dim)
self.fc2 = nn.Linear(feed_forward_dim, hidden_dim)
def forward(self, x):
x = nn.functional.relu(self.fc1(x))
x = self.fc2(x)
return x
```
以上代码实现了一个基本的Transformer模型,包括Encoder、Decoder、MultiheadAttention等组件。你可以根据自己的需求进行修改和扩展。