class Transformer(nn.Module): def __init__(self, vocab_size: int, max_seq_len: int, embed_dim: int, hidden_dim: int, n_layer: int, n_head: int, ff_dim: int, embed_drop: float, hidden_drop: float): super().__init__() self.tok_embedding = nn.Embedding(vocab_size, embed_dim) self.pos_embedding = nn.Embedding(max_seq_len, embed_dim) layer = nn.TransformerEncoderLayer( d_model=hidden_dim, nhead=n_head, dim_feedforward=ff_dim, dropout=hidden_drop) self.encoder = nn.TransformerEncoder(layer, num_layers=n_layer) self.embed_dropout = nn.Dropout(embed_drop) self.linear1 = nn.Linear(embed_dim, hidden_dim) self.linear2 = nn.Linear(hidden_dim, embed_dim) def encode(self, x, mask): x = x.transpose(0, 1) x = self.encoder(x, src_key_padding_mask=mask) x = x.transpose(0, 1) return x
时间: 2024-04-26 12:24:19 浏览: 176
这是一段使用 PyTorch 实现的 Transformer 模型的代码,用于自然语言处理任务中的序列建模,例如文本分类、机器翻译等。
该模型的输入是一个词汇表大小为 `vocab_size`,最大序列长度为 `max_seq_len` 的词嵌入(embedding)矩阵,其中每个词嵌入的维度为 `embed_dim`。模型使用了 `n_layer` 层 TransformerEncoderLayer,每个 EncoderLayer 中包含了 `n_head` 个注意力头(self-attention)。每个 EncoderLayer 的隐藏层大小为 `hidden_dim`,Feedforward 层的大小为 `ff_dim`,并在每个 EncoderLayer 后应用了一个 `hidden_drop` 的 Dropout。在模型的输入层和第一个 EncoderLayer 之间,使用了一个 `embed_drop` 的 Dropout。
在 forward 方法中,输入的 `x` 是一个形状为 `(batch_size, seq_len)` 的整数张量,表示一个批次中的多个序列。`mask` 是一个形状为 `(batch_size, seq_len)` 的布尔型张量,用于指示哪些位置是填充值,需要被屏蔽。在 encode 方法中,模型首先将输入的 `x` 转置为 `(seq_len, batch_size)` 的形状,然后将其输入到 TransformerEncoder 中进行编码。最后,将编码结果再次转置为 `(batch_size, seq_len)` 的形状并返回。