PyTorch学习笔记：针对一个网络的权重初始化方法

# 针对一个网络的权重初始化方法
import torch
import torch.nn as nn


## 建立一个测试网络
class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.hidden = nn.Sequential(
            nn.Linear(100, 100),
            nn.ReLU(),
            nn.Linear(100, 50),
            nn.ReLU(),
        )
        self.cla = nn.Linear(50, 10)

    # 定义网络的前向传播路径
    def forward(self, x):
        x = self.conv1(x)
        x = x.view(x.shape[0], -1)
        x = self.hidden(x)
        output = self.cla(x)
        return output


def init_weights(m):
    '''
    对不同类型层的参数使用不同的方法进行初始化
    :param m: 网络的某一层
    :return:
    '''
    # 如果是卷积层
    if type(m) == nn.Conv2d:
        torch.nn.init.normal_(m.weight, mean=0, std=0.5)
        # UserWarning: nn.init.normal is now deprecated in favor of nn.init.normal_.
    # 如果是全连接层
    if type(m) == nn.Linear:
        torch.nn.init.uniform_(m.weight, a=-0.1, b=0.1)
        m.bias.data.fill_(0.01)


if __name__ == '__main__':
    # 输出网络结构
    testnet = TestNet()
    print(testnet)

    # 使用网络的apply方法进行权重初始化
    torch.manual_seed(13)
    testnet.apply(init_weights)

    print(testnet.conv1.weight)  # 输出参数