论文:PolyGen: An Autoregressive Generative Model of 3D Meshes
首先阅读transformer铺垫知识《Torch中Transformer
的中文注释》。
以下为Encoder
部分,很简单,小学生都会:
from typing import Dict, Optional, Tuple
import pdbimport torch
import torch.nn as nn
from torch.nn import MultiheadAttention, Linear, Dropout, LayerNorm, ReLU, Parameter
import pytorch_lightning as plfrom .utils import get_clonesclass PolygenDecoderLayer(nn.TransformerDecoderLayer):"""根据Vaswani等人2017年的描述,这是一个解码器模块。它使用了遮蔽自注意力和非遮蔽跨注意力来处理序列化上下文模块。实现了缓存机制以加快解码速度。缓存的作用是存储键值对,这样在解码器的每次前向传递中就不必重新生成这些键值对。参数:d_model: 嵌入向量的大小。nhead: 多头注意力机制中的头数。dim_feedforward: 全连接层的大小。dropout: 在每个连接层后ReLU激活函数之后应用的dropout比率。re_zero: 如果为True,则使用alpha比例因子对残差进行零初始化。构造函数 __init__初始化参数包括d_model(嵌入向量的大小)、nhead(多头注意力机制中的头数)、dim_feedforward(全连接层的大小)、dropout(dropout比率)和re_zero(是否使用re_zero技术)。创建了多个多头注意力模块,包括自注意力(self.self_attn)和跨注意力(self.multihead_attn)。定义了前馈网络所需的线性层和激活函数。设置了层归一化(LayerNorm)和Dropout层。如果re_zero被设置为True,则初始化三个可学习的参数alpha、beta和gamma,用于调整残差连接中的权重,这有助于训练深度网络时的稳定性。"""def __init__(self,d_model: int = 256,nhead: int = 4,dim_feedforward: int = 1024,dropout: float = 0.2,re_zero: bool = True,) -> None:super(PolygenDecoderLayer, self).__init__(d_model, nhead, dim_feedforward=dim_feedforward, dropout=dropout)# 初始化多头自注意力和多头跨注意力模块self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)# 初始化前馈网络self.linear1 = Linear(d_model, dim_feedforward)self.dropout = Dropout(dropout)self.linear2 = Linear(dim_feedforward, d_model)# 初始化层归一化self.norm1 = LayerNorm(d_model)self.norm2 = LayerNorm(d_model)self.norm3 = LayerNorm(d_model)# 初始化Dropoutself.dropout1 = Dropout(dropout)self.dropout2 = Dropout(dropout)self.dropout3 = Dropout(dropout)# 初始化激活函数self.activation = ReLU()# 初始化re_zero参数self.re_zero = re_zeroself.alpha = Parameter(data=torch.Tensor([0.0]))self.beta = Parameter(data=torch.Tensor([0.0]))self.gamma = Parameter(data=torch.Tensor([0.0]))def forward(self,tgt: torch.Tensor,memory: Optional[torch.Tensor] = None,tgt_mask: Optional[torch.Tensor] = None,memory_mask: Optional[torch.Tensor] = None,tgt_key_padding_mask: Optional[torch.Tensor] = None,memory_key_padding_mask: Optional[torch.Tensor] = None,cache: Optional[Dict[str, torch.Tensor]] = None,) -> torch.Tensor:"""解码器层的前向传播方法。参数:tgt: 输入序列的张量,形状为 [sequence_length, batch_size, embed_size]。(目标序列)memory: 编码器最后一层的序列张量,形状为 [source_sequence_length, batch_size, embed_size]。(来自编码器的输出)tgt_mask: 目标序列的掩码张量,形状为 [sequence_length, sequence_length]。memory_mask: 存储序列的掩码张量,形状为 [sequence_length, source_sequence_length]。tgt_key_padding_mask: 目标序列的键填充掩码张量,形状为 [batch_size, sequence_length]。memory_key_padding_mask: 存储序列的键填充掩码张量,形状为 [batch_size, source_sequence_length]。cache: 用于快速解码的缓存,格式为 {'k': torch.Tensor, 'v': torch.Tensor} 的字典。返回:输出张量,形状为 [sequence_length, batch_size, embed_size],是一个解码器层的前向循环结果。处理流程:接收tgt(目标序列)、memory(来自编码器的输出)、各种掩码和cache作为输入参数。如果提供了cache,会将当前的tgt追加到保存的键值对中,以便于后续的解码步骤能够利用历史信息。使用自注意力层处理输入tgt,考虑到了tgt_mask和tgt_key_padding_mask。如果re_zero开启,自注意力的输出乘以alpha,控制每层的贡献度。接着使用跨注意力层处理memory,同样考虑掩码,并使用beta调整输出。最后通过前馈网络,其中使用gamma调整前馈网络的输出。每个子层的输出都与输入进行了残差连接,并经过Dropout。NOTE:Cache technique.被用在Transformer模型的解码器层;cache中已经保存了之前解码步骤中的键和值张量。在解码新词时,这些已有的键和值张量会被与当前的tgt(目标序列)张量拼接在一起,形成更新后的键和值张量,然后将这些更新后的键和值张量再次存入cache中,供下一次解码使用"""# 处理缓存以实现快速解码if cache is not None:saved_key = cache["k"]saved_value = cache["v"]key = cache["k"] = torch.cat(tensors=[saved_key, tgt], axis=0)value = cache["v"] = torch.cat(tensors=[saved_value, tgt], axis=0)else:key = tgtvalue = tgt# 应用自注意力层tgt2 = self.norm1(tgt)tgt2 = self.self_attn(tgt, key, value, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]# 如果启用了re_zero,应用比例因子if self.re_zero:tgt2 = tgt2 * self.alpha # 输出与一个比例因子(self.alpha)相乘,这有助于控制每层的贡献,从而稳定训练过程。tgt = tgt + self.dropout1(tgt2)# 应用跨注意力层if memory is not None:tgt2 = self.norm2(tgt)tgt2 = self.multihead_attn(tgt,memory.float(),memory.float(),attn_mask=memory_mask,key_padding_mask=memory_key_padding_mask,)[0]if self.re_zero:tgt2 = tgt2 * self.betatgt2 = self.dropout2(tgt2)tgt = tgt + tgt2# 应用前馈网络tgt2 = self.norm3(tgt)tgt2 = self.linear1(tgt2)tgt2 = self.activation(tgt2)tgt2 = self.linear2(tgt2)if self.re_zero:tgt2 = tgt2 * self.gammatgt2 = self.dropout(tgt2)tgt = tgt + tgt2return tgtclass PolygenDecoder(pl.LightningModule):"""Polygen的解码器,修改自Pytorch的Transformer解码器实现,增加了缓存机制用于加速解码过程。"""def __init__(self,decoder_layer: nn.TransformerDecoderLayer,num_layers: int,norm: Optional[nn.Module] = None,) -> None:"""Polygen解码器的初始化函数。参数:decoder_layer: nn.TransformerDecoderLayer 类型的 Pytorch 模块,指定解码器层的具体实现。num_layers: 解码器中层的数量。norm: 在解码器的前向传播方法结束时应用的层归一化类型。初始化 (__init__ 方法):构造函数通过 get_clones 函数初始化 layers 属性,创建给定 decoder_layer 的多个副本,每个副本都是 nn.TransformerDecoderLayer 的独立实例。同时设置 num_layers 和 norm 属性。"""super(PolygenDecoder, self).__init__()self.layers = get_clones(decoder_layer, num_layers) # 获取多个解码器层的副本self.num_layers = num_layers # 解码器层数self.norm = norm # 层归一化def forward(self,tgt: torch.Tensor,memory: Optional[torch.Tensor] = None,tgt_mask: Optional[torch.Tensor] = None,memory_mask: Optional[torch.Tensor] = None,tgt_key_padding_mask: Optional[torch.Tensor] = None,memory_key_padding_mask: Optional[torch.Tensor] = None,cache: Optional[Tuple] = None,) -> torch.Tensor:"""解码器层的前向传播方法。参数:tgt: 输入序列的张量,形状为 [sequence_length, batch_size, embed_size]。memory: 编码器最后一层的序列张量,形状为 [source_sequence_length, batch_size, embed_size]。tgt_mask: 目标序列的掩码张量,形状为 [sequence_length, sequence_length]。memory_mask: 存储序列的掩码张量,形状为 [sequence_length, source_sequence_length]。tgt_key_padding_mask: 忽略目标序列中指定填充元素的张量,形状为 [batch_size, sequence_length]。memory_key_padding_mask: 忽略存储序列中指定填充元素的张量,形状为 [batch_size, source_sequence_length]。cache: 用于快速解码的缓存,格式为每个层的 {'k': torch.Tensor, 'v': torch.Tensor} 字典组成的列表。返回值:输出张量,形状为 [sequence_length, batch_size, embed_size],是所有解码器层的前向传播结果。处理流程:遍历每层解码器,将当前的 output 张量应用到这一层上。如果提供了 cache,则获取当前层的缓存 (layer_cache) 并将 tgt_mask 设置为 None。这是因为缓存实际上通过存储过去的键值对来内部处理掩码,使得显式掩码变得不必要。若没有缓存,layer_cache 被设置为 None,表明不应使用任何缓存值。对于每一层,通过将 output 以及 memory 和相关掩码(以及可用的缓存)传入当前层,更新 output 张量。如果设置了 norm,在所有解码器层完成之后,对 output 应用层归一化。"""output = tgtfor i, mod in enumerate(self.layers):if cache is not None:layer_cache = cache[i] # 使用缓存加速解码tgt_mask = None # 使用缓存时,目标掩码置空else:layer_cache = None # 不使用缓存output = mod( # 将当前层的输入、编码器的输出、相关的掩码和缓存(如果有)作为参数,计算出新的输出张量。output,memory,tgt_mask=tgt_mask,memory_mask=memory_mask,tgt_key_padding_mask=tgt_key_padding_mask,memory_key_padding_mask=memory_key_padding_mask,cache=layer_cache,)if self.norm is not None:output = self.norm(output) # 应用层归一化return outputclass TransformerDecoder(pl.LightningModule):def __init__(self,hidden_size: int = 256,fc_size: int = 1024,num_heads: int = 4,layer_norm: bool = True,num_layers: int = 8,dropout_rate: float = 0.2,) -> None:"""Transformer解码器,结合了PolygenDecoderLayer和PolygenDecoder。参数:hidden_size: 嵌入向量的大小。fc_size: 全连接层的大小。num_heads: 多头注意力头的数量。layer_norm: 是否使用层归一化。num_layers: 解码器中的层数。dropout_rate: ReLU之后立即应用的dropout比率。"""super(TransformerDecoder, self).__init__()# 初始化解码器结构self.hidden_size = hidden_sizeself.num_layers = num_layersself.decoder = PolygenDecoder(PolygenDecoderLayer(d_model=hidden_size,nhead=num_heads,dim_feedforward=fc_size,dropout=dropout_rate,),num_layers=num_layers,norm=LayerNorm(hidden_size), # 使用层归一化)def initialize_cache(self, batch_size) -> Dict[str, torch.Tensor]:"""初始化用于快速解码的缓存。参数:batch_size: 输入批次的大小。返回:cache: 包含特定解码器层键值对的字典列表。"""# 初始化k和v张量k = torch.zeros([0, batch_size, self.hidden_size], device=self.device)v = torch.zeros([0, batch_size, self.hidden_size], device=self.device)cache = [{"k": k, "v": v} for _ in range(self.num_layers)]return cachedef generate_square_subsequent_mask(self, sz: int) -> torch.Tensor:"""为输入序列生成一个目标掩码。为序列自注意力机制生成一个掩码,使得每个元素只能关注它之前的元素(包括自己),这样就防止了信息的未来泄露。生成一个下三角掩码矩阵,用于确保在序列自注意力机制中,每个位置只能看到它之前的元素(包括自己)参数:sz: 输入序列的长度。返回:mask: 形状为[sequence_length, sequence_length]的下三角矩阵。流程:假设我们有一个序列长度sz为5,即我们的序列长度是51. 下三角阵mask = (torch.triu(torch.ones(sz, sz, device=self.device)) == 1).transpose(0, 1)创建一个5x5的全1张量;然后将其传递给torch.triu函数,以获取一个上三角矩阵;转为下三角阵。2. 转换类型并应用掩码布尔型的下三角矩阵转为浮点型; 使用masked_fill函数将False的位置(即矩阵中非下三角的部分)填充为负无穷大(float('-inf')),而将True的位置(即下三角部分)填充为0tensor([[ 0., -inf, -inf, -inf, -inf],[ 0., 0., -inf, -inf, -inf],[ 0., 0., 0., -inf, -inf],[ 0., 0., 0., 0., -inf],[ 0., 0., 0., 0., 0.]])"""# 生成掩码,确保自注意力只关注不晚于当前位置的位置# 生成一个下三角掩码矩阵,用于确保序列中的每个位置只能注意到该位置之前(包含自身)的位置。mask = (torch.triu(torch.ones(sz, sz, device=self.device)) == 1).transpose(0, 1)mask = mask.float().masked_fill(mask == 0, float("-inf")).masked_fill(mask == 1, float(0.0))return maskdef forward(self,inputs: torch.Tensor,sequential_context_embeddings: Optional[torch.Tensor] = None,cache: Optional[tuple] = None,) -> torch.Tensor:"""Transformer解码器的前向传播方法。输入序列(inputs)、上下文嵌入(sequential_context_embeddings)和缓存(cache)参数:inputs: 形状为[sequence_length, batch_size, embed_size]的张量,代表输入序列。cache: 一个字典列表,每个字典格式为{'k': torch.Tensor, 'v': torch.Tensor},代表各解码器层的缓存。返回:out: 解码器所有层的前向传播结果,形状为[sequence_length, batch_size, embed_size]的张量。NOTE:tgt: A Tensor of shape [sequence_length, batch_size, embed_size]. Represents the input sequence.memory: A Tensor of shape [source_sequence_length, batch_size, embed_size]. Represents the sequence from the last layer of the encoder.tgt_mask: A Tensor of shape [sequence_length, sequence_length]. The mask for the target sequence. 下三角阵,控制模型不能开到`未来`的知识cache: A list of dictionaries in the following format: {'k': torch.Tensor, 'v': torch.Tensor}. Each dictionary in the list represents the cache at the respective decoder layer."""sz = inputs.shape[0] # 输入序列长度mask = self.generate_square_subsequent_mask(sz) # 生成自注意力掩码out = self.decoder(inputs, memory=sequential_context_embeddings, tgt_mask=mask, cache=cache) # 执行前向传播return out