pytorch 权重组成分析

一般网络使用pytorch训练，权重保存在pth中。

读取pth使用  state_dict = torch.load('path/*.pth', map_location='cpu')

得到的state_dict是一个字典，key为层的name，value为weight。假设block1为由conv2d，BN组成的，那么卷积对应的key为block1.0.weight， block1.0.bias，BN层参数较多，包括running_mean ，running_var，gamma，beta，eps， num_batches_tracked，其前缀为block1.1。

for key in state_dict.keys():
    print(key, state_dict[key].shape)  # 打印key和value的shape
    del state_dict[key] # 删除

假设网络定义为net ，直接通过net也可以打印key，param为随机的，然后通过load_state_dict加在权重。

for weights_name, param in net.named_parameters():
    print(weights_name, param)
    param.data = torch.zeros(size)  # 对权重就行赋值
net.load_state_dict(state_dict)
torch.save(state_dict, "./ckpts/new.pth") # 保存新权重

btw，named_modules()可以显示net的模块，比如net由block0， block1， block2组成

for name, module in net.named_modules():
    print(name, module)  #  name 为block0,1,2;module为网络定义