PyTorch深度学习实战 |层归一化层和FeedForward

import torch
from torch import nn
# 定义张量x
# x是一个2×4大小的张量，代表了2个样本，每个样本有4个特征
# 当使用层归一化对x进行计算时，会对x的每一行的4个特征进行归一化
x = torch.tensor([[1.0, 2.0, 3.0, 4.0],
                  [1.0, 1.0, 1.0, 1.0]])
# 定义层归一化实例
# 传入的参数4，表示层归一化的输入维度是4
# 也就是对应某一层的4个神经元的输出
layer_norm = nn.LayerNorm(4)
# 计算层归一化的结果
output = layer_norm(x)
print(output)

import torch
import torch.nn as nn
# 当使用层归一化计算序列数据时，层归一化只会作用在最后的特征维度D上
# 其中批量大小B与序列长度Q，这两个维度都不会影响层归一化的计算
# 定义输入数据 x，形状为 (B, Q, D) -> (2, 3, 4)
# 2个样本，每个样本3个单词，每个单词用4个维度表示
x = torch.tensor([
    # 第1个样本：
    [
        [1.0, 2.0, 3.0, 4.0],  # 第1个单词的4个维度
        [1.0, 1.0, 1.0, 1.0],  # 第2个单词的4个维度
        [1.0, 1.0, 1.0, 1.0]   # 第3个单词的4个维度
    ],
    # 第2个样本：
    [
        [1.0, 1.0, 1.0, 1.0],  # 第1个单词的4个维度
        [1.0, 1.0, 1.0, 1.0],  # 第2个单词的4个维度
        [1.0, 1.0, 1.0, 1.0]   # 第3个单词的4个维度
    ]
])
# 初始化层归一化层
layer_norm = nn.LayerNorm(4)
# x输入至层归一化layer_norm中
output = layer_norm(x)
# 输出结果
print(output)

import torch
import torch.nn as nn
import torch.nn.functional as F
class FeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int = None, activation: str = "gelu", dropout: float = 0.1):
        """
        Args:
            d_model: 输入/输出维度（即 Transformer 的隐藏层维度）
            d_ff:    中间层维度（默认扩展为 4*d_model）
            activation: 激活函数，支持 "gelu" 或 "relu"
            dropout:  Dropout 概率
        """
        super().__init__()
        d_ff = d_ff or 4 * d_model  # 默认扩展比例为4倍
        # 定义两个线性层
        self.linear1 = nn.Linear(d_model, d_ff)  # 扩展维度
        self.linear2 = nn.Linear(d_ff, d_model)  # 压缩回原维度
        
        # 激活函数选择
        self.activation = F.gelu if activation == "gelu" else F.relu
        
        # Dropout 层（可选）
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        # 参数初始化（关键！）
        self._init_weights()
    def _init_weights(self):
        # 使用 He/Kaiming 初始化（适合 ReLU）
        nn.init.kaiming_normal_(self.linear1.weight, nonlinearity='relu')
        nn.init.zeros_(self.linear1.bias)
        
        # 缩小输出层的初始化范围
        nn.init.xavier_normal_(self.linear2.weight, gain=0.02)
        nn.init.zeros_(self.linear2.bias)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        输入形状: (batch_size, seq_len, d_model)
        输出形状: (batch_size, seq_len, d_model)
        """
        x = self.linear1(x)       # (batch, seq, d_ff)
        x = self.activation(x)    # 非线性激活
        x = self.dropout(x)       # 可选 Dropout
        x = self.linear2(x)       # (batch, seq, d_model)
        return x