在PyTorch中,可以使用torch.nn.Transformer
類來調用Transformer模型。以下是一個使用Transformer模型的示例代碼:
import torch
import torch.nn as nn
# 定義Transformer模型
class TransformerModel(nn.Module):
def __init__(self, input_dim, hidden_dim, num_layers, num_heads):
super(TransformerModel, self).__init__()
self.transformer = nn.Transformer(
d_model=input_dim,
nhead=num_heads,
num_encoder_layers=num_layers,
num_decoder_layers=num_layers,
dim_feedforward=hidden_dim
)
def forward(self, src, tgt):
output = self.transformer(src, tgt)
return output
# 創建Transformer模型實例
input_dim = 512
hidden_dim = 2048
num_layers = 6
num_heads = 8
model = TransformerModel(input_dim, hidden_dim, num_layers, num_heads)
# 準備輸入數據
batch_size = 16
src_seq_len = 10
tgt_seq_len = 5
src = torch.randn(batch_size, src_seq_len, input_dim)
tgt = torch.randn(batch_size, tgt_seq_len, input_dim)
# 前向傳播
output = model(src, tgt)
在這個示例中,我們首先定義了一個繼承自nn.Module
的自定義Transformer模型類TransformerModel
。在__init__
方法中,我們使用nn.Transformer
類來創建一個Transformer模型,并指定輸入維度、隱藏層維度、編碼器和解碼器的層數,以及注意力頭數。在forward
方法中,我們將輸入數據傳入Transformer模型進行前向傳播,并返回輸出。
然后,我們創建了一個Transformer模型實例,并準備了輸入數據。最后,我們通過調用模型的forward
方法來進行前向傳播,并得到輸出結果。