transformer使用示例

#加载数据train_data, test_data, vocab, pos_vocab = load_treebank()

def load_treebank():    # 需要翻墙下载，可以自行设置代码    nltk.set_proxy('http://192.168.0.28:1080')    # 如果没有的话那么则会下载，否则忽略    nltk.download('treebank')    from nltk.corpus import treebank    sents, postags = zip(*(zip(*sent) for sent in treebank.tagged_sents()))    vocab = Vocab.build(sents, reserved_tokens=["<pad>"])    tag_vocab = Vocab.build(postags)    train_data = [(vocab.convert_tokens_to_ids(sentence), tag_vocab.convert_tokens_to_ids(tags)) for sentence, tags in zip(sents[:3000], postags[:3000])]    test_data = [(vocab.convert_tokens_to_ids(sentence), tag_vocab.convert_tokens_to_ids(tags)) for sentence, tags in zip(sents[3000:], postags[3000:])]    return train_data, test_data, vocab, tag_vocab

>>> train_data[0]>>> Out[1]: ([2, 3, 4, 5, 6, 7, 4, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], [1, 1, 2, 3, 4, 5, 2, 6, 7, 8, 9, 10, 8, 5, 9, 1, 3, 11])

# 这个函数就是将其变成等长，填充使用<pad>，至于是0还是1还是其他值并不重要，因为还有mask~def collate_fn(examples):    lengths = torch.tensor([len(ex[0]) for ex in examples])    inputs = [torch.tensor(ex[0]) for ex in examples]    targets = [torch.tensor(ex[1]) for ex in examples]    inputs = pad_sequence(inputs, batch_first=True, padding_value=vocab["<pad>"])    targets = pad_sequence(targets, batch_first=True, padding_value=vocab["<pad>"])    return inputs, lengths, targets, inputs != vocab["<pad>"]

class PositionalEncoding(nn.Module):    def __init__(self, d_model, dropout=0.1, max_len=512):        super(PositionalEncoding, self).__init__()        pe = torch.zeros(max_len, d_model)        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))        pe[:, 0::2] = torch.sin(position * div_term)        pe[:, 1::2] = torch.cos(position * div_term)        pe = pe.unsqueeze(0).transpose(0, 1)        self.register_buffer('pe', pe)    def forward(self, x):        x = x + self.pe[:x.size(0), :]        return xclass Transformer(nn.Module):    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_class,                 dim_feedforward=512, num_head=2, num_layers=2, dropout=0.1, max_len=512, activation: str = "relu"):        super(Transformer, self).__init__()        # 词嵌入层        self.embedding_dim = embedding_dim        self.embeddings = nn.Embedding(vocab_size, embedding_dim)        self.position_embedding = PositionalEncoding(embedding_dim, dropout, max_len)        # 编码层：使用Transformer        encoder_layer = nn.TransformerEncoderLayer(hidden_dim, num_head, dim_feedforward, dropout, activation)        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)        # 输出层        self.output = nn.Linear(hidden_dim, num_class)    def forward(self, inputs, lengths):        inputs = torch.transpose(inputs, 0, 1)        hidden_states = self.embeddings(inputs)        hidden_states = self.position_embedding(hidden_states)        attention_mask = length_to_mask(lengths) == False        hidden_states = self.transformer(hidden_states, src_key_padding_mask=attention_mask).transpose(0, 1)        logits = self.output(hidden_states)        log_probs = F.log_softmax(logits, dim=-1)        return log_probs