init_bert_weights代码

以下是基于PyTorch实现的init_bert_weights代码，用于初始化BERT模型的权重参数：

import torch.nn as nn
import torch.nn.init as init

def init_bert_weights(module):
    """BERT权重初始化"""
    if isinstance(module, (nn.Linear, nn.Embedding)):
        # 模型中的Linear和Embedding层的权重初始化
        module.weight.data.normal_(mean=0.0, std=0.02)
        if module.bias is not None:
            module.bias.data.zero_()
    elif isinstance(module, nn.LayerNorm):
        # 模型中的LayerNorm层的权重初始化
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)
    if isinstance(module, nn.Linear) and module.bias is not None:
        # 模型中的Linear层的偏置项初始化
        module.bias.data.zero_()

# 使用示例：
# 初始化BERT模型的权重参数
model = BertModel.from_pretrained('bert-base-uncased')
model.apply(init_bert_weights)

此代码会对BERT模型的所有权重进行初始化，其中：

- Linear和Embedding层的权重使用均值为0，标准差为0.02的正态分布进行初始化，偏置项初始化为0；
- LayerNorm层的偏置项初始化为0，权重初始化为1。

这些参数初始化方式与原论文保持一致。