Transformer PyTorch Pseudo Code

transformer-outline.py

import torch.nn as nn


class Transformer(nn.Module):
    def __init__(self, ...):
        super().__init__()
        ...
        self.encoder_layers = nn.ModuleList(...)
        self.decoder_layers = nn.ModuleList(...)
        ...

    def forward(self, x, y):
        ...
        
        encoder_result = x
        for layer in self.encoder_layers:
            encoder_result = layer(encoder_result, x_mask)

        decoder_result = y
        for layer in self.decoder_layers:
            decoder_result = layer(decoder_result, encoder_result, x_mask, y_mask)

        output = ...
        return output

loss = nn.CrossEntropyLoss(0)
optimizer = optim.AdamW(...)
  
for batch in data_loader:
    optimizer.zero_grad()
    output = transformer(x_data, y_data[:, :-1])
    loss = loss(...)
    loss.backward()
    optimizer.step()