【深度学习:理论篇】--完整代码实现Transformer
本文实现了一个基于PyTorch的Transformer模型，包含完整的编码器-解码器架构。主要模块包括：1）多头自注意力机制(SelfAttention)，实现Q/K/V计算和掩码处理；2）Transformer块(TransformerBlock)，整合自注意力、前馈网络和残差连接；3）编码器(Encoder)，包含词嵌入、位置编码和多层Transformer块；4）解码器(Decoder)，
·
目录
1.注意力机制
X是输入，W是线性变换矩阵
class SelfAttention(nn.Module):
def __init__(self, embed_size, heads):
super(SelfAttention, self).__init__()
self.embed_size = embed_size # 输入特征维度
self.heads = heads # 注意力头的数量
self.head_dim = embed_size // heads # 每个头的维度
# 确保embed_size能被heads整除
assert (self.head_dim * heads == embed_size), "Embed size needs to be div by heads"
# 定义Q、K、V的线性变换层（无偏置）
self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
# 多头注意力拼接后的输出层
self.fc_out = nn.Linear(heads * self.head_dim, embed_size)
def forward(self, values, keys, query, mask):
N = query.shape[0] # 批大小（样本数量）
value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1] # 获取序列长度
# 将输入拆分为多头（reshape成多头形式），reshape成(N, seq_len, heads, head_dim)
values = values.reshape(N, value_len, self.heads, self.head_dim)
keys = keys.reshape(N, key_len, self.heads, self.head_dim)
queries = query.reshape(N, query_len, self.heads, self.head_dim)
# 对Q、K、V进行线性变换
values = self.values(values)
keys = self.keys(keys)
queries = self.queries(queries)
# 计算注意力分数（query和key的点积）
energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
# 输出形状：(N, heads, query_len, key_len)
# 应用mask（将mask为0的位置填充负无穷）
if mask is not None:
energy = energy.masked_fill(mask == 0, float("-1e20"))
# 计算注意力权重（softmax归一化）
attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)
# 注意力权重与value相乘得到输出
out = torch.einsum("nhql,nlhd->nqhd", [attention, values])
out = out.reshape(N, query_len, self.heads * self.head_dim) # 拼接多头结果
# 通过输出层进行线性变换
out = self.fc_out(out)
return out
2.Transformer单块
完整流程、残差链接、
前馈神经网络：
实现整个流程
class TransformerBlock(nn.Module):
def __init__(self, embed_size, heads, dropout, forward_expansion):
#参数说明
# embed_size: 输入特征维度
# heads: 注意力头的数量
# dropout: Dropout层的概率
# forward_expansion: 前馈神经网络的扩展倍数
super(TransformerBlock, self).__init__()
# 初始化自注意力机制层
self.attention = SelfAttention(embed_size, heads)
# 初始化两个LayerNorm层，用于残差连接后的归一化
self.norm1 = nn.LayerNorm(embed_size)
self.norm2 = nn.LayerNorm(embed_size)
# 前馈神经网络（FFN）部分
self.feed_forward = nn.Sequential(
nn.Linear(embed_size, forward_expansion * embed_size), # 扩展维度
nn.ReLU(), # 激活函数
nn.Linear(forward_expansion * embed_size, embed_size) # 恢复原始维度
)
# Dropout层，用于防止过拟合
self.dropout = nn.Dropout(dropout)
def forward(self, value, key, query, mask):
# 1. 计算自注意力
attention = self.attention(value, key, query, mask)
# 2. 残差连接 + LayerNorm + Dropout
# 注意：原始Transformer论文是先Norm再残差，这里实现的是先残差再Norm的变体
x = self.dropout(self.norm1(attention + query))
# 3. 前馈神经网络
forward = self.feed_forward(x)
# 4. 第二次残差连接 + LayerNorm + Dropout
out = self.dropout(self.norm2(forward + x))
return out3.编码器（encoder）
class Encoder(nn.Module):
def __init__(self,
src_vocab_size, # (int) 源语言词汇表大小，决定词嵌入矩阵的行数
embed_size, # (int) 词嵌入维度，每个词/位置将被映射到的向量维度
num_layers, # (int) Transformer编码器堆叠的层数
heads, # (int) 多头注意力机制的头数
device, # (torch.device) 计算设备 (CPU/GPU)
forward_expansion, # (int) 前馈网络隐藏层的扩展倍数，决定FFN中间层维度
dropout, # (float) Dropout概率，用于正则化
max_length # (int) 最大序列长度，决定位置嵌入矩阵的行数
):
super(Encoder, self).__init__()
# 初始化参数
self.embed_size = embed_size # (int) 保存词嵌入维度
self.device = device # (torch.device) 保存计算设备
# 词嵌入层：将词汇索引映射为embed_size维向量
# 参数:
# - num_embeddings: (int) 词汇表大小(src_vocab_size)
# - embedding_dim: (int) 嵌入维度(embed_size)
self.word_embedding = nn.Embedding(src_vocab_size, embed_size)
# 位置嵌入层：将位置索引映射为embed_size维向量
# 参数:
# - num_embeddings: (int) 最大序列长度(max_length)
# - embedding_dim: (int) 嵌入维度(embed_size)
self.position_embedding = nn.Embedding(max_length, embed_size)
# 创建num_layers个TransformerBlock组成的编码器堆叠
# 每个TransformerBlock参数:
# - embed_size: (int) 输入/输出维度
# - heads: (int) 注意力头数
# - dropout: (float) Dropout概率
# - forward_expansion: (int) FFN扩展倍数
self.layers = nn.ModuleList(
[
TransformerBlock(
embed_size,
heads,
dropout=dropout,
forward_expansion=forward_expansion
)
for _ in range(num_layers)] # 循环num_layers次
)
# Dropout层
# 参数:
# - p: (float) Dropout概率(dropout)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask):
"""
前向传播
参数:
x: (torch.Tensor) 输入张量，形状为(N, seq_length)，
N是batch大小，seq_length是序列长度，包含词汇索引
mask: (torch.Tensor) 掩码张量，用于屏蔽无效位置(如padding位置)
返回:
(torch.Tensor) 编码后的输出张量，形状为(N, seq_length, embed_size)，
包含序列中每个位置的上下文感知表示
"""
# 获取batch大小(N)和序列长度(seq_length)
N, seq_lengh = x.shape
# 生成位置索引: [0, 1, 2, ..., seq_length-1]
# 并扩展到batch维度 -> (N, seq_length)
positions = torch.arange(0, seq_lengh).expand(N, seq_lengh).to(self.device)
# 1. 词嵌入: (N, seq_length) -> (N, seq_length, embed_size)
# 2. 位置嵌入: (N, seq_length) -> (N, seq_length, embed_size)
# 3. 相加后进行dropout
out = self.dropout(self.word_embedding(x) + self.position_embedding(positions))
# 逐层通过Transformer编码器
for layer in self.layers:
# 每层的输入作为Q,K,V (自注意力机制)
# mask用于屏蔽无效位置
out = layer(out, out, out, mask) # Q=K=V=out
# 返回编码后的表示 (N, seq_length, embed_size)
return out4.解码器（decoder）
class Decoder(nn.Module):
def __init__(self,
trg_vocab_size,
embed_size,
num_layers,
heads,
forward_expansion,
dropout,
device,
max_length):
"""
解码器初始化
参数:
trg_vocab_size: 目标词汇表大小
embed_size: 嵌入维度
num_layers: 解码器层数
heads: 注意力头数量
forward_expansion: 前向扩展因子
dropout: dropout概率
device: 计算设备
max_length: 最大序列长度(用于位置编码)
"""
super(Decoder, self).__init__()
self.device = device
# 词嵌入层
self.word_embedding = nn.Embedding(trg_vocab_size, embed_size)
# 位置嵌入层
self.position_embedding = nn.Embedding(max_length, embed_size)
# 创建多层解码器块
self.layers = nn.ModuleList(
[DecoderBlock(embed_size, heads, forward_expansion, dropout, device)
for _ in range(num_layers)]
)
# 输出全连接层(将嵌入维度映射回词汇表大小)
self.fc_out = nn.Linear(embed_size, trg_vocab_size)
# Dropout层
self.dropout = nn.Dropout(dropout)
def forward(self, x, enc_out, src_mask, trg_mask):
"""
前向传播
参数:
x: 输入序列(目标语言)
enc_out: 编码器输出(用于key和value)
src_mask: 源序列掩码
trg_mask: 目标序列掩码
返回:
解码后的输出(词汇表概率分布)
"""
# 获取batch大小和序列长度
N, seq_length = x.shape
# 创建位置索引(0到seq_length-1)
positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)
# 词嵌入 + 位置嵌入 + dropout
x = self.dropout((self.word_embedding(x) + self.position_embedding(positions)))
# 通过所有解码器层
for layer in self.layers:
x = layer(x, enc_out, enc_out, src_mask, trg_mask)
# 通过全连接层得到输出(词汇表概率分布)
out = self.fc_out(x)
return out5.完整Transform
class Transformer(nn.Module):
def __init__(self,
src_vocab_size, # 源语言词汇表大小
trg_vocab_size, # 目标语言词汇表大小
src_pad_idx, # 源语言padding索引
trg_pad_idx, # 目标语言padding索引
embed_size=256, # 嵌入维度(默认256)
num_layers=6, # 编码器/解码器层数(默认6)
forward_expansion=4, # 前向扩展因子(默认4)
heads=8, # 多头注意力头数(默认8)
dropout=0, # dropout率(默认0)
device="cuda", # 计算设备(默认cuda)
max_length=100 # 最大序列长度(默认100)
):
super(Transformer, self).__init__()
# 初始化编码器
self.encoder = Encoder(
src_vocab_size,
embed_size,
num_layers,
heads,
device,
forward_expansion,
dropout,
max_length
)
# 初始化解码器
self.decoder = Decoder(
trg_vocab_size,
embed_size,
num_layers,
heads,
forward_expansion,
dropout,
device,
max_length
)
# 保存padding索引和设备信息
self.src_pad_idx = src_pad_idx
self.trg_pad_idx = trg_pad_idx
self.device = device
def make_src_mask(self, src):
"""
创建源语言掩码
参数:
src: 源语言输入序列
返回:
src_mask: 源语言掩码张量(N, 1, 1, src_len)
其中padding位置为False，其他为True
"""
# 创建布尔掩码(padding位置为False)
src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)
return src_mask.to(self.device)
def make_trg_mask(self, trg):
"""
创建目标语言掩码
参数:
trg: 目标语言输入序列
返回:
trg_mask: 目标语言掩码张量(N, 1, trg_len, trg_len)
包含1)防止看到未来信息的三角掩码
2) padding位置的掩码
"""
N, trg_len = trg.shape
# 创建下三角矩阵(防止看到未来信息)
trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(
N, 1, trg_len, trg_len
)
return trg_mask.to(self.device)
def forward(self, src, trg):
"""
前向传播
参数:
src: 源语言输入序列
trg: 目标语言输入序列
返回:
out: 模型输出(目标语言词汇表上的概率分布)
"""
# 创建源语言掩码
src_mask = self.make_src_mask(src)
# 创建目标语言掩码
trg_mask = self.make_trg_mask(trg)
# 编码器处理源语言
enc_src = self.encoder(src, src_mask)
# 解码器生成目标语言
out = self.decoder(trg, enc_src, src_mask, trg_mask)
return out
6.完整可运行代码
import torch
import torch.nn as nn
class SelfAttention(nn.Module):
def __init__(self, embed_size, heads):
super(SelfAttention, self).__init__()
self.embed_size = embed_size # 输入特征维度
self.heads = heads # 注意力头的数量
self.head_dim = embed_size // heads # 每个头的维度
# 确保embed_size能被heads整除
assert (self.head_dim * heads == embed_size), "Embed size needs to be div by heads"
# 定义Q、K、V的线性变换层（无偏置）
self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
# 多头注意力拼接后的输出层
self.fc_out = nn.Linear(heads * self.head_dim, embed_size)
def forward(self, values, keys, query, mask):
N = query.shape[0] # 批大小（样本数量）
value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1] # 获取序列长度
# 将输入拆分为多头（reshape成多头形式），reshape成(N, seq_len, heads, head_dim)
values = values.reshape(N, value_len, self.heads, self.head_dim)
keys = keys.reshape(N, key_len, self.heads, self.head_dim)
queries = query.reshape(N, query_len, self.heads, self.head_dim)
# 对Q、K、V进行线性变换
values = self.values(values)
keys = self.keys(keys)
queries = self.queries(queries)
# 计算注意力分数（query和key的点积）
energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
# 输出形状：(N, heads, query_len, key_len)
# 应用mask（将mask为0的位置填充负无穷）
if mask is not None:
energy = energy.masked_fill(mask == 0, float("-1e20"))
# 计算注意力权重（softmax归一化）
attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)
# 注意力权重与value相乘得到输出
out = torch.einsum("nhql,nlhd->nqhd", [attention, values])
out = out.reshape(N, query_len, self.heads * self.head_dim) # 拼接多头结果
# 通过输出层进行线性变换
out = self.fc_out(out)
return out
class TransformerBlock(nn.Module):
def __init__(self, embed_size, heads, dropout, forward_expansion):
#参数说明
# embed_size: 输入特征维度
# heads: 注意力头的数量
# dropout: Dropout层的概率
# forward_expansion: 前馈神经网络的扩展倍数
super(TransformerBlock, self).__init__()
# 初始化自注意力机制层
self.attention = SelfAttention(embed_size, heads)
# 初始化两个LayerNorm层，用于残差连接后的归一化
self.norm1 = nn.LayerNorm(embed_size)
self.norm2 = nn.LayerNorm(embed_size)
# 前馈神经网络（FFN）部分
self.feed_forward = nn.Sequential(
nn.Linear(embed_size, forward_expansion * embed_size), # 扩展维度
nn.ReLU(), # 激活函数
nn.Linear(forward_expansion * embed_size, embed_size) # 恢复原始维度
)
# Dropout层，用于防止过拟合
self.dropout = nn.Dropout(dropout)
def forward(self, value, key, query, mask):
# 1. 计算自注意力
attention = self.attention(value, key, query, mask)
# 2. 残差连接 + LayerNorm + Dropout
# 注意：原始Transformer论文是先Norm再残差，这里实现的是先残差再Norm的变体
x = self.dropout(self.norm1(attention + query))
# 3. 前馈神经网络
forward = self.feed_forward(x)
# 4. 第二次残差连接 + LayerNorm + Dropout
out = self.dropout(self.norm2(forward + x))
return out
class Encoder(nn.Module):
def __init__(self,
src_vocab_size, # (int) 源语言词汇表大小，决定词嵌入矩阵的行数
embed_size, # (int) 词嵌入维度，每个词/位置将被映射到的向量维度
num_layers, # (int) Transformer编码器堆叠的层数
heads, # (int) 多头注意力机制的头数
device, # (torch.device) 计算设备 (CPU/GPU)
forward_expansion, # (int) 前馈网络隐藏层的扩展倍数，决定FFN中间层维度
dropout, # (float) Dropout概率，用于正则化
max_length # (int) 最大序列长度，决定位置嵌入矩阵的行数
):
super(Encoder, self).__init__()
# 初始化参数
self.embed_size = embed_size # (int) 保存词嵌入维度
self.device = device # (torch.device) 保存计算设备
# 词嵌入层：将词汇索引映射为embed_size维向量
# 参数:
# - num_embeddings: (int) 词汇表大小(src_vocab_size)
# - embedding_dim: (int) 嵌入维度(embed_size)
self.word_embedding = nn.Embedding(src_vocab_size, embed_size)
# 位置嵌入层：将位置索引映射为embed_size维向量
# 参数:
# - num_embeddings: (int) 最大序列长度(max_length)
# - embedding_dim: (int) 嵌入维度(embed_size)
self.position_embedding = nn.Embedding(max_length, embed_size)
# 创建num_layers个TransformerBlock组成的编码器堆叠
# 每个TransformerBlock参数:
# - embed_size: (int) 输入/输出维度
# - heads: (int) 注意力头数
# - dropout: (float) Dropout概率
# - forward_expansion: (int) FFN扩展倍数
self.layers = nn.ModuleList(
[
TransformerBlock(
embed_size,
heads,
dropout=dropout,
forward_expansion=forward_expansion
)
for _ in range(num_layers)] # 循环num_layers次
)
# Dropout层
# 参数:
# - p: (float) Dropout概率(dropout)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask):
"""
前向传播
参数:
x: (torch.Tensor) 输入张量，形状为(N, seq_length)，
N是batch大小，seq_length是序列长度，包含词汇索引
mask: (torch.Tensor) 掩码张量，用于屏蔽无效位置(如padding位置)
返回:
(torch.Tensor) 编码后的输出张量，形状为(N, seq_length, embed_size)，
包含序列中每个位置的上下文感知表示
"""
# 获取batch大小(N)和序列长度(seq_length)
N, seq_lengh = x.shape
# 生成位置索引: [0, 1, 2, ..., seq_length-1]
# 并扩展到batch维度 -> (N, seq_length)
positions = torch.arange(0, seq_lengh).expand(N, seq_lengh).to(self.device)
# 1. 词嵌入: (N, seq_length) -> (N, seq_length, embed_size)
# 2. 位置嵌入: (N, seq_length) -> (N, seq_length, embed_size)
# 3. 相加后进行dropout
out = self.dropout(self.word_embedding(x) + self.position_embedding(positions))
# 逐层通过Transformer编码器
for layer in self.layers:
# 每层的输入作为Q,K,V (自注意力机制)
# mask用于屏蔽无效位置
out = layer(out, out, out, mask) # Q=K=V=out
# 返回编码后的表示 (N, seq_length, embed_size)
return out
class DecoderBlock(nn.Module):
def __init__(self, embed_size, heads, forward_expansion, dropout, device):
"""
解码器块初始化
参数:
embed_size: 嵌入维度大小
heads: 注意力头的数量
forward_expansion: 前向扩展因子(用于TransformerBlock)
dropout: dropout概率
device: 计算设备(CPU/GPU)
"""
super(DecoderBlock, self).__init__()
# 自注意力机制
self.attention = SelfAttention(embed_size, heads)
# 层归一化
self.norm = nn.LayerNorm(embed_size)
# Transformer块(包含多头注意力、前馈网络等)
self.transformer_block = TransformerBlock(
embed_size, heads, dropout, forward_expansion
)
# Dropout层
self.dropout = nn.Dropout(dropout)
def forward(self, x, value, key, src_mask, trg_mask):
"""
前向传播
参数:
x: 输入张量
value: 来自编码器的value
key: 来自编码器的key
src_mask: 源序列掩码(用于屏蔽padding等)
trg_mask: 目标序列掩码(用于防止未来信息泄露)
返回:
解码后的输出
"""
# 自注意力计算(使用目标序列掩码)
attention = self.attention(x, x, x, trg_mask)
# 残差连接 + 层归一化 + dropout
query = self.dropout(self.norm(attention + x))
# 通过Transformer块(使用编码器的key和value)
out = self.transformer_block(value, key, query, src_mask)
return out
class Decoder(nn.Module):
def __init__(self,
trg_vocab_size,
embed_size,
num_layers,
heads,
forward_expansion,
dropout,
device,
max_length):
"""
解码器初始化
参数:
trg_vocab_size: 目标词汇表大小
embed_size: 嵌入维度
num_layers: 解码器层数
heads: 注意力头数量
forward_expansion: 前向扩展因子
dropout: dropout概率
device: 计算设备
max_length: 最大序列长度(用于位置编码)
"""
super(Decoder, self).__init__()
self.device = device
# 词嵌入层
self.word_embedding = nn.Embedding(trg_vocab_size, embed_size)
# 位置嵌入层
self.position_embedding = nn.Embedding(max_length, embed_size)
# 创建多层解码器块
self.layers = nn.ModuleList(
[DecoderBlock(embed_size, heads, forward_expansion, dropout, device)
for _ in range(num_layers)]
)
# 输出全连接层(将嵌入维度映射回词汇表大小)
self.fc_out = nn.Linear(embed_size, trg_vocab_size)
# Dropout层
self.dropout = nn.Dropout(dropout)
def forward(self, x, enc_out, src_mask, trg_mask):
"""
前向传播
参数:
x: 输入序列(目标语言)
enc_out: 编码器输出(用于key和value)
src_mask: 源序列掩码
trg_mask: 目标序列掩码
返回:
解码后的输出(词汇表概率分布)
"""
# 获取batch大小和序列长度
N, seq_length = x.shape
# 创建位置索引(0到seq_length-1)
positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device)
# 词嵌入 + 位置嵌入 + dropout
x = self.dropout((self.word_embedding(x) + self.position_embedding(positions)))
# 通过所有解码器层
for layer in self.layers:
x = layer(x, enc_out, enc_out, src_mask, trg_mask)
# 通过全连接层得到输出(词汇表概率分布)
out = self.fc_out(x)
return out
class Transformer(nn.Module):
def __init__(self,
src_vocab_size, # 源语言词汇表大小
trg_vocab_size, # 目标语言词汇表大小
src_pad_idx, # 源语言padding索引
trg_pad_idx, # 目标语言padding索引
embed_size=256, # 嵌入维度(默认256)
num_layers=6, # 编码器/解码器层数(默认6)
forward_expansion=4, # 前向扩展因子(默认4)
heads=8, # 多头注意力头数(默认8)
dropout=0, # dropout率(默认0)
device="cuda", # 计算设备(默认cuda)
max_length=100 # 最大序列长度(默认100)
):
super(Transformer, self).__init__()
# 初始化编码器
self.encoder = Encoder(
src_vocab_size,
embed_size,
num_layers,
heads,
device,
forward_expansion,
dropout,
max_length
)
# 初始化解码器
self.decoder = Decoder(
trg_vocab_size,
embed_size,
num_layers,
heads,
forward_expansion,
dropout,
device,
max_length
)
# 保存padding索引和设备信息
self.src_pad_idx = src_pad_idx
self.trg_pad_idx = trg_pad_idx
self.device = device
def make_src_mask(self, src):
"""
创建源语言掩码
参数:
src: 源语言输入序列
返回:
src_mask: 源语言掩码张量(N, 1, 1, src_len)
其中padding位置为False，其他为True
"""
# 创建布尔掩码(padding位置为False)
src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)
return src_mask.to(self.device)
def make_trg_mask(self, trg):
"""
创建目标语言掩码
参数:
trg: 目标语言输入序列
返回:
trg_mask: 目标语言掩码张量(N, 1, trg_len, trg_len)
包含1)防止看到未来信息的三角掩码
2) padding位置的掩码
"""
N, trg_len = trg.shape
# 创建下三角矩阵(防止看到未来信息)
trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(
N, 1, trg_len, trg_len
)
return trg_mask.to(self.device)
def forward(self, src, trg):
"""
前向传播
参数:
src: 源语言输入序列
trg: 目标语言输入序列
返回:
out: 模型输出(目标语言词汇表上的概率分布)
"""
# 创建源语言掩码
src_mask = self.make_src_mask(src)
# 创建目标语言掩码
trg_mask = self.make_trg_mask(trg)
# 编码器处理源语言
enc_src = self.encoder(src, src_mask)
# 解码器生成目标语言
out = self.decoder(trg, enc_src, src_mask, trg_mask)
return out
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
x = torch.tensor([[1, 5, 6, 4, 3, 9, 5, 2, 0], [1, 8, 7, 3, 4, 5, 6, 7, 2]]).to(
device
)
trg = torch.tensor([[1, 7, 4, 3, 5, 9, 2, 0], [1, 5, 6, 2, 4, 7, 6, 2]]).to(device)
src_pad_idx = 0#源语言padding索引
trg_pad_idx = 0#目标语言padding索引
src_vocab_size = 10#源语言词汇表大小
trg_vocab_size = 10#目标语言词汇表大小
model = Transformer(src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, device=device).to(device)
out = model(x, trg[:, :-1])
print(out.shape)
魔乐社区（Modelers.cn) 是一个中立、公益的人工智能社区，提供人工智能工具、模型、数据的托管、展示与应用协同服务，为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作，由全产业链共同建设、共同运营、共同享有，推动国产AI生态繁荣发展。
更多推荐
13
0- 0
已为社区贡献27条内容
所有评论(0)