如何在PyTorch中初始化权重

技术背景

在深度学习中，神经网络的权重初始化是一个重要的步骤，合适的权重初始化方法可以加速模型的收敛速度，提高模型的性能。PyTorch作为一个广泛使用的深度学习框架，提供了多种权重初始化的方法。

实现步骤

单一层的权重初始化

对于单个层的权重初始化，可以使用torch.nn.init中的函数。例如，使用Xavier均匀分布初始化卷积层的权重：

import torch
import torch.nn as nn

conv1 = nn.Conv2d(...)
torch.nn.init.xavier_uniform_(conv1.weight)

也可以直接修改weight.data和bias.data来初始化权重和偏置：

conv1.weight.data.fill_(0.01)
conv1.bias.data.fill_(0.01)

nn.Sequential或自定义nn.Module的权重初始化

可以使用torch.nn.Module.apply方法递归地初始化整个nn.Module的权重。例如：

import torch.nn as nn

def init_weights(m):
    if isinstance(m, nn.Linear):
        torch.nn.init.xavier_uniform_(m.weight)
        m.bias.data.fill_(0.01)

net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)

核心代码

不同分布的权重初始化

import torch
import torch.nn as nn
import numpy as np

# 均匀分布初始化
def weights_init_uniform(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        m.weight.data.uniform_(0.0, 1.0)
        m.bias.data.fill_(0)

# 通用规则初始化
def weights_init_uniform_rule(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        n = m.in_features
        y = 1.0/np.sqrt(n)
        m.weight.data.uniform_(-y, y)
        m.bias.data.fill_(0)

# 正态分布初始化
def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        y = m.in_features
        m.weight.data.normal_(0.0, 1/np.sqrt(y))
        m.bias.data.fill_(0)

model_uniform = Net()
model_uniform.apply(weights_init_uniform)

model_rule = Net()
model_rule.apply(weights_init_uniform_rule)

model_normal = Net()
model_normal.apply(weights_init_normal)

手动设置权重

import torch
import torch.nn as nn

input = torch.ones((8, 8))
d = nn.Linear(8, 8, bias=False)
d.weight.data = torch.full((8, 8), 0.5)
output = d(input)

最佳实践

• 根据激活函数选择初始化方法：对于ReLU激活函数，Kaiming He初始化通常是一个不错的选择；对于Sigmoid或Tanh激活函数，Xavier初始化可能更合适。
• 避免所有权重初始化为相同值：如果所有权重初始化为0或1，会导致所有神经元的输出相同，使得模型难以学习。
• 参考默认初始化：PyTorch在很多层中已经有了默认的初始化方法，如Linear层使用Kaiming He初始化，可以在不需要特殊初始化时直接使用。

常见问题

不指定初始化时的默认初始化方法是什么？

对于Linear层，默认使用Kaiming He初始化：

import math
import torch.nn.init as init

init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
    fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
    bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
    init.uniform_(self.bias, -bound, bound)

为什么不能将所有权重初始化为0或1？

如果所有权重相同，每个层的所有神经元都会产生相同的输出，这使得很难确定应该调整哪些权重，导致模型难以学习。

如何使用正态分布进行权重初始化？

可以使用torch.nn.init.normal_函数：

torch.nn.init.normal_(tensor, mean=0, std=1)