transformer 序列分类

Transformer 模型可以用于序列分类任务，其中输入是一个序列，输出是一个标签或分类。下面是一个简单的 Transformer 序列分类模型的示例代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import TransformerEncoder, TransformerEncoderLayer

class TransformerClassifier(nn.Module):
    def __init__(self, input_dim, output_dim, n_layers, hidden_dim, n_heads, dropout):
        super(TransformerClassifier, self).__init__()

        self.model_type = 'Transformer'
        self.src_mask = None
        self.pos_encoder = PositionalEncoding(hidden_dim, dropout)
        encoder_layers = TransformerEncoderLayer(hidden_dim, n_heads, hidden_dim, dropout)
        self.transformer_encoder = TransformerEncoder(encoder_layers, n_layers)
        self.encoder = nn.Linear(input_dim, hidden_dim)
        self.decoder = nn.Linear(hidden_dim, output_dim)

    def forward(self, src):
        if self.src_mask is None or self.src_mask.size(0) != len(src):
            mask = self._generate_square_subsequent_mask(len(src)).to(src.device)
            self.src_mask = mask

        src = self.encoder(src)
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src, self.src_mask)
        output = self.decoder(output.mean(dim=0))
        return F.log_softmax(output, dim=1)

    def _generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask.to(src.device)

该模型将输入序列传递给一个线性编码器，然后通过位置编码器进行编码。接下来，通过 TransformerEncoderLayer 和 TransformerEncoder 进行多层编码。最后，通过线性解码器将输出映射到标签空间，并通过 log_softmax 函数进行归一化。