在模型中使用Kaiming初始化的代码

在PyTorch中，可以使用`torch.nn.init`模块中的`kaiming_normal_`或`kaiming_uniform_`函数来进行Kaiming初始化。这两个函数分别可以初始化权重为均值为0，方差为 $\frac{2}{n}$ 的高斯分布或均匀分布，其中n是前一层的神经元数量。

例如，对于一个全连接层，可以使用以下代码来使用Kaiming初始化：

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 30)
        self.fc3 = nn.Linear(30, 2)

        # 使用Kaiming初始化
        init.kaiming_normal_(self.fc1.weight)
        init.kaiming_normal_(self.fc2.weight)
        init.kaiming_normal_(self.fc3.weight)
        
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

在以上代码中，`init.kaiming_normal_`函数被用来初始化`fc1`、`fc2`和`fc3`的权重，其中`kaiming_normal_`表示使用高斯分布进行初始化。如果要使用均匀分布进行初始化，可以使用`kaiming_uniform_`函数。