记录

pytorch中模型的各参数初始化

方法

1、self.modules() 方法

用于获取自定义类中所有模块、层。通过遍历所有层，实现逐一初始化。

class MyNet(nn.Module):
    def __init__(self):
    	# 在这里实例化各种层（Conv2D、BN、Attention ... ）
		for m in self.modules():	# 获取所有层
            if isinstance(m, nn.Conv2d):	# 对卷积层初始化，标准分布
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d): 	# 对BN层初始化，方差为1，偏差为0
            	m.weight.data.fill_(1)
                m.bias.data.zero_()

2、kaiming初始化

def init_conv_kaiming(self,layer):
    if isinstance(layer, nn.Conv2d):
        nn.init.kaiming_uniform_(layer.weight, a=1)

        if layer.bias is not None:
            nn.init.constant_(layer.bias, 0)