PyTorch torch.nn.Transformer 関数
PyTorch torch.nn リファレンスマニュアルtorch.nn.TransformerPyTorch の完全な Transformer モデルです。
エンコーダーとデコーダーが含まれており、シーケンス・ツー・シーケンスのタスクに使用できます。
関数定義
torch.nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, activation='gelu', batch_first=True)
使用例
例 1: 基本的な使い方
実例
import torch
import torch.nn as nn
# Transformer を作成
transformer = nn.Transformer(d_model=512, nhead=8, batch_first=True)
# エンコーダー入力
src = torch.randn(10, 32, 512) # (seq, batch, d_model)
# デコーダー入力
tgt = torch.randn(20, 32, 512)
output = transformer(src, tgt)
print("出力形状:", output.shape)
import torch.nn as nn
# Transformer を作成
transformer = nn.Transformer(d_model=512, nhead=8, batch_first=True)
# エンコーダー入力
src = torch.randn(10, 32, 512) # (seq, batch, d_model)
# デコーダー入力
tgt = torch.randn(20, 32, 512)
output = transformer(src, tgt)
print("出力形状:", output.shape)
例 2: 簡単な翻訳モデル
実例
import torch
import torch.nn as nn
class TransformerMT(nn.Module):
def __init__(self, vocab_size, d_model=256, nhead=4):
super(TransformerMT, self).__init__()
self.d_model = d_model
self.embedding = nn.Embedding(vocab_size, d_model)
self.transformer = nn.Transformer(d_model, nhead, batch_first=True)
self.fc = nn.Linear(d_model, vocab_size)
def forward(self, src, tgt):
src = self.embedding(src) * (self.d_model ** 0.5)
tgt = self.embedding(tgt) * (self.d_model ** 0.5)
out = self.transformer(src, tgt)
return self.fc(out)
model = TransformerMT(10000)
src = torch.randint(0, 10000, (32, 50))
tgt = torch.randint(0, 10000, (32, 40))
output = model(src, tgt)
print("出力形状:", output.shape)
import torch.nn as nn
class TransformerMT(nn.Module):
def __init__(self, vocab_size, d_model=256, nhead=4):
super(TransformerMT, self).__init__()
self.d_model = d_model
self.embedding = nn.Embedding(vocab_size, d_model)
self.transformer = nn.Transformer(d_model, nhead, batch_first=True)
self.fc = nn.Linear(d_model, vocab_size)
def forward(self, src, tgt):
src = self.embedding(src) * (self.d_model ** 0.5)
tgt = self.embedding(tgt) * (self.d_model ** 0.5)
out = self.transformer(src, tgt)
return self.fc(out)
model = TransformerMT(10000)
src = torch.randint(0, 10000, (32, 50))
tgt = torch.randint(0, 10000, (32, 40))
output = model(src, tgt)
print("出力形状:", output.shape)
使用シーン
- 機械翻訳
- テキスト生成
- シーケンス・ツー・シーケンス
ヒント：batch_first=True の場合、入力形状は (batch, seq, d_model) です。
その他の拡張機能