pytorch模型加载测试_pytorch加载模型和初始化权重

1. 初始化权重

对网络中的某一层进行初始化

self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)

init.xavier_uniform(self.conv1.weight)

init.constant(self.conv1.bias, 0.1)

对网络的整体进行初始化:

def weights_init(m):

classname=m.__class__.__name__

if classname.find('Conv') != -1:

xavier(m.weight.data)

xavier(m.bias.data)

net = Net()#构建网络

net.apply(weights_init) #apply函数会递归地搜索网络内的所有module并把参数表示的函数应用到所有的module上。

#对所有的Conv层都初始化权重.

权重初始化

for m in model.modules():

if isinstance(m,nn.Conv2d):

nn.init.normal(m.weight.data)

nn.init.xavier_normal(m.weight.data)

nn.init.kaiming_normal(m.weight.data)#卷积层参数初始化

m.bias.data.fill_(0)

elif isinstance(m,nn.Linear):

m.weight.data.normal_()#全连接层参数初始化

2. 加载模型

pretrained_dict = ...

<