一个基本的Transformer分类器的示例代码

这是一个基本的Transformer分类器的示例代码：

import torch
import torch.nn as nn

class TransformerClassifier(nn.Module):
  def __init__(self, num_classes, num_tokens, hidden_size=512, num_attention_heads=8, num_layers=6):
    super(TransformerClassifier, self).__init__()
    self.transformer = nn.Transformer(
      d_model=hidden_size, 
      nhead=num_attention_heads, 
      num_encoder_layers=num_layers, 
      num_decoder_layers=num_layers
    )
    self.classifier = nn.Linear(hidden_size, num_classes)
    self.init_weights()
  
  def init_weights(self):
    # Initialize the weights of the linear layer
    nn.init.xavier_uniform_(self.classifier.weight)
    nn.init.zeros_(self.classifier.bias)
  
  def forward(self, input_ids, attention_mask=None):
    # Pass the input through the transformer model
    output = self.transformer(input_ids, attention_mask=attention_mask)[0]
    # Take the mean of the output along the sequence dimension
    mean_output = output.mean(dim=1)
    # Pass the mean through the linear layer to get the logits
    logits = self.classifier(mean_output)
    return logits

# Instantiate the model with num_classes=2 and num_tokens=20000
model = TransformerClassifier(num_classes=2, num_tokens=20000)

# Define the input
input_ids = torch.LongTensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
attention_mask = torch.LongTensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 1]])

# Get the logits
logits = model(input_ids, attention_mask=attention_mask)
print(logits)

这个模型接受两个输入：input_ids 和 attention_mask。input_ids 是一个形状为 (batch_size, sequence_length) 的张量，表示输入文本的单词编号。attention_mask 是一个形状为 (batch_size, sequence_length) 的张量，表示每个位置是否需要考虑。

在这个模型中，我们使用了 nn.Transformer 模型来