用pytorch构造GELU激活函数

本文介绍如何在Python 2及低版本PyTorch中,由于缺少GELU模块,作者通过自定义类实现了GELU激活函数,重点讨论了两种近似计算方法并比较其与nn.GELU的差异。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

因为一些原因必须使用python2+低版本的pytorch来跑代码,其中就遇到了低版本pytorch没有封装GELU激活函数的情况,所以在网上把该函数的源码找到了,自己构造了一个基于nn.Module的版本:

class GELU(nn.Module):#zengen
    def __init__(self):
        super(GELU, self).__init__()
    def forward(self, x):
        #return 0.5*x*(1+torch.tanh(np.sqrt(2/np.pi)*(x+0.044715*torch.pow(x,3))))
        return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

其中第一个注释掉了的return是gelu原文中提出的一种近似计算版本,而下面那一个是pytorch中封装的版本,实测第一种会和nn.GELU计算的值略有差异,因此使用第二种。

# 4. PreNormTransformerEncoderLayer class PreNormTransformerEncoderLayer(nn.Module): def __init__(self, d_model: int, nhead: int, dim_feedforward: int = 2048, dropout: float = 0.1, batch_first: bool = True, activation: nn.Module = nn.GELU(), bias: bool = False): super().__init__() self.self_attn = nn.MultiheadAttention( embed_dim=d_model, num_heads=nhead, dropout=dropout, batch_first=batch_first, bias=bias ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) # 激活函数处理 if isinstance(activation, str): activation = activation.lower() if activation == "relu": self.activation = nn.ReLU() elif activation == "gelu": self.activation = nn.GELU() else: raise ValueError(f"Unsupported activation: {activation}") else: self.activation = activation # 前馈网络 self.feedforward = nn.Sequential( nn.Linear(d_model, dim_feedforward, bias=bias), self.activation, nn.Dropout(dropout), nn.Linear(dim_feedforward, d_model, bias=bias) ) self._reset_parameters() def _reset_parameters(self): # 自注意力参数初始化 nn.init.xavier_uniform_(self.self_attn.in_proj_weight) if self.self_attn.in_proj_bias is not None: nn.init.constant_(self.self_attn.in_proj_bias, 0.) nn.init.xavier_uniform_(self.self_attn.out_proj.weight) if self.self_attn.out_proj.bias is not None: nn.init.constant_(self.self_attn.out_proj.bias, 0.) # 前馈网络改进初始化 nn.init.kaiming_uniform_( # 改为Kaiming初始化 self.feedforward[0].weight, nonlinearity='gelu' if isinstance(self.activation, nn.GELU) else 'relu' ) if self.feedforward[0].bias is not None: nn.init.constant_(self.feedforward[0].bias, 0.) nn.init.xavier_uniform_(self.feedforward[-1].weight) # 最后一层保持Xavier if self.feedforward[-1].bias is not None: nn.init.constant_(self.feedforward[-1].bias, 0.) def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor: # 复用LayerNorm计算结果 normed_src = self.norm1(src) attn_output = self.self_attn( normed_src, normed_src, normed_src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask, need_weights=False )[0] src = src + self.dropout1(attn_output) # Feedforward normed_src2 = self.norm2(src) ff_output = self.feedforward(normed_src2) src = src + self.dropout2(ff_output) return src
03-19
<think>嗯,我现在需要理解这个PreNormTransformerEncoderLayer类的代码。首先,我应该先通读一遍代码,看看各个部分是怎么组成的。这个类继承自nn.Module,所以是一个PyTorch的模块。构造函数里的参数包括d_model,nhead,dim_feedforward,dropout等等,这些都是Transformer中常见的参数。 首先,self.self_attn是多头注意力层,参数看起来都是标准的,比如embed_dim是d_model,头数是nhead,dropout,batch_first和bias。这里bias参数在初始化的时候被设为False,可能作者觉得去掉偏置项可以提升模型效果或者减少参数? 接下来是两个LayerNorm层,norm1和norm2,然后是dropout层。激活函数部分,用户可以通过传入字符串或者直接传模块来选择,比如"relu"或"gelu"。默认是GELU,这可能因为GELU在Transformer中表现更好? 前馈网络feedforward是一个序列,包含两个线性层,中间有激活函数和dropout。这里结构应该是先升维到dim_feedforward,再降回d_model,标准的Transformer前馈结构。 然后有_reset_parameters方法,用来初始化参数。这里注意到自注意力的参数用了Xavier均匀初始化,而前馈网络的第一层用了Kaiming初始化,最后一层还是Xavier。这可能是因为中间使用GELU或ReLU,Kaiming更适合这类激活函数,而输出层保持Xavier有助于稳定输出? 在前向传播部分,输入src经过norm1之后,作为自注意力的query、key和value。这里用的是Pre-Norm结构,即先做归一化再进行注意力计算。然后加上残差连接和dropout。接着,再做一次norm2和前馈网络,同样残差连接加dropout。所以整个流程是:PreNorm → Attention → Add & Dropout → PreNorm → FFN → Add & Dropout。这跟传统的Post-Norm不同,Pre-Norm可能有助于训练稳定性,特别是在深层网络中。 这里需要注意的是,前向传播中,两次都用了不同的norm层,norm1和norm2。而每次归一化后的结果被用于计算,然后加上原始输入,这符合残差结构的设计。例如,attn_output是基于normed_src计算的,然后加到src上,再经过第二个norm和FFN。 参数初始化部分,自注意力的in_proj_weight是Xavier均匀初始化,out_proj也是。前馈网络的第一层用Kaiming,考虑到激活函数的不同,比如GELU可能需要不同的初始化方式。而最后一层用Xavier可能为了保持输出的稳定性。这部分可能需要进一步思考为什么这样选择,是否有相关论文支持。 另外,代码中的bias参数在初始化时被设置为False,那么在Linear层和MultiheadAttention中是否使用bias呢?例如,在前馈网络的Linear层中,是否设置了bias?在代码中,前馈网络的两个Linear层都传入了bias=bias,所以如果构造函数中的bias是False,那么这些层就不会有偏置项。 还有,src_mask和src_key_padding_mask的作用,在调用self_attn的时候传入了这两个mask参数。src_mask应该是用于屏蔽某些注意力权重(比如防止看到未来的词),而src_key_padding_mask是用于屏蔽padding的位置。这些参数在训练时如何处理可能需要看具体的使用场景。 在激活函数处理部分,代码允许传入字符串或者模块实例。如果是字符串,会转换为小写,然后选择对应的激活函数。默认是GELU,这可能更符合现代Transformer的设计,比如Vision Transformer等。 残差连接的部分,每次都是src = src + dropout(output),这样保留了原始的输入,有助于梯度流动。这里两次残差连接,分别对应自注意力和前馈网络。 初始化方法中,前馈网络的第一层使用Kaiming均匀初始化,根据激活函数是否为GELU或ReLU来选择nonlinearity参数。而最后一层使用Xavier均匀初始化,这可能因为输出需要保持与输入相当的尺度,避免梯度爆炸或消失。 总的来说,这个PreNormTransformerEncoderLayer实现了Pre-Normalization的Transformer编码器层,与原始Transformer的Post-Norm不同。这样的设计可能在训练时更稳定,尤其是在较深的网络中。每个子层(自注意力和前馈)之前都有LayerNorm,然后加上残差连接,这种结构可能有助于信息的流动。 不过,需要确认的是,Pre-Norm和Post-Norm的区别。Post-Norm是在子层之后进行归一化,而Pre-Norm是在子层之前。例如,原始Transformer是Post-Norm,即在自注意力之后做Add & Norm,而这里是先做Norm,再自注意力,然后Add。这样的改变可能影响模型的学习动态,但具体效果可能需要实验验证。 此外,参数初始化的策略也值得注意,不同的层使用不同的初始化方法,这可能对模型训练的收敛性和效果有影响。比如,Kaiming初始化更适合有ReLU/GELU激活函数的层,而Xavier可能更适合线性变换。 最后,代码中的dropout被应用在自注意力和前馈网络的输出之后,这有助于防止过拟合。同时,LayerNorm的参数在初始化时应该使用的是默认的PyTorch初始化方式,即gamma为1,beta为0,除非有特别设置。 可能的问题点:在_reset_parameters方法中,是否覆盖了PyTorch默认的LayerNorm初始化?因为代码里没有对norm1和norm2进行特别的初始化,所以它们的参数会保持PyTorch默认的,即gamma为1,beta为0。这可能没问题,因为LayerNorm通常不需要特别的初始化。 总结,这个类实现了一个基于Pre-Normalization的Transformer编码器层,包含了多头自注意力机制和前馈网络,使用了特定的参数初始化策略,并支持不同的激活函数。适用于需要更稳定训练的深层Transformer模型。</think>以下是对`PreNormTransformerEncoderLayer`的逐步解析: --- ### **1. 模块功能定位** 这是一个基于**Pre-Normalization**结构的Transformer编码器层实现,与传统Transformer的Post-Norm不同。主要特点包括: - **先归一化再计算子层(自注意力/前馈)** - **残差连接 + Dropout** - **可配置激活函数与初始化策略** --- ### **2. 核心组件分解** #### **2.1 自注意力层** ```python self.self_attn = nn.MultiheadAttention(...) ``` - **参数说明**: - `embed_dim=d_model`: 输入/输出维度 - `num_heads=nhead`: 多头注意力头数 - `bias=False`: 禁用偏置项(减少参数量) #### **2.2 归一化层** ```python self.norm1 = nn.LayerNorm(d_model) # 自注意力前的归一化 self.norm2 = nn.LayerNorm(d_model) # 前馈前的归一化 ``` - **作用**:稳定训练过程,加速收敛 #### **2.3 前馈网络** ```python self.feedforward = nn.Sequential( nn.Linear(d_model → dim_feedforward), activation, # 默认GELU Dropout, nn.Linear(dim_feedforward → d_model) ) ``` - **结构特点**:先扩展维度(通常4倍)再压缩回原维度 --- ### **3. 初始化策略** #### **3.1 自注意力参数** ```python # 权重:Xavier均匀分布 nn.init.xavier_uniform_(self.self_attn.in_proj_weight) # 偏置:归零(若存在) ``` - **目的**:保持注意力权重的平稳初始化 #### **3.2 前馈网络参数** ```python # 第一层:Kaiming初始化(适配GELU/ReLU) nn.init.kaiming_uniform_(self.feedforward[0].weight, nonlinearity='gelu') # 最后一层:Xavier初始化 nn.init.xavier_uniform_(self.feedforward[-1].weight) ``` - **设计意图**: - 中间层适配激活函数的分布特性 - 输出层维持输入输出维度平衡 --- ### **4. 前向传播流程** ```python def forward(...): # Pre-Norm自注意力 normed_src = self.norm1(src) attn_output = self.self_attn(normed_src, ...) src = src + self.dropout1(attn_output) # 残差连接 # Pre-Norm前馈 normed_src2 = self.norm2(src) ff_output = self.feedforward(normed_src2) src = src + self.dropout2(ff_output) # 残差连接 return src ``` - **关键步骤**: 1. **输入归一化** → 计算自注意力 → 残差连接 2. **再次归一化** → 前馈计算 → 残差连接 --- ### **5. Pre-Norm vs Post-Norm** | 特性 | Pre-Norm (本文) | Post-Norm (原始Transformer) | |---------------------|------------------------------------------|---------------------------------------| | **归一化位置** | 子层计算前 | 子层计算后 | | **训练稳定性** | 更稳定(梯度直接传播) | 可能需精细调参 | | **深层表现** | 适合深层网络 | 需辅助手段(如残差缩放) | | **典型应用** | T5、GPT-3等现代模型 | 原始Transformer、BERT | --- ### **6. 设计亮点总结** 1. **结构优化**: - Pre-Norm提升梯度流动效率 - 双残差连接防止信息丢失 2. **参数配置**: - 自适应激活函数选择(GELU默认) - 分层的初始化策略 3. **工程细节**: - 屏蔽偏置项减少计算量 - 完整的Mask支持(padding/sequence mask) --- ### **7. 适用场景建议** - **长序列建模**:Pre-Norm对长程依赖更友好 - **深层Transformer**:如12层以上的编码器 - **需快速收敛的任务**:如小规模数据集 如需进一步了解具体训练技巧或调参方法,可提供更多上下文继续探讨。
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值