Fonction torch.nn.Transformer de PyTorch

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

torch.nn.TransformerC'est le modèle Transformer complet dans PyTorch.

Il contient un encodeur et un décodeur, et peut être utilisé pour des tâches de séquence à séquence.

Définition de fonction

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)

Exemples d'utilisation

Exemple 1 : Utilisation de base

Exemple

import torch
import torch.nn as nn

# Créer un Transformer
transformer = nn.Transformer(d_model=512, nhead=8, batch_first=True)

# Entrée de l'encodeur
src = torch.randn(10, 32, 512)  # (seq, batch, d_model)
# Entrée du décodeur
tgt = torch.randn(20, 32, 512)

output = transformer(src, tgt)

print("Forme de sortie :", output.shape)

Exemple 2 : Modèle de traduction simple

Exemple

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("Forme de sortie :", output.shape)

Cas d'utilisation

  • Traduction automatique
  • Génération de texte
  • Séquence-à-séquence

Remarque : lorsque batch_first=True, la forme de l'entrée est (batch, seq, d_model).

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

Autres extensions