[PyTorch]论文pytorch复现中遇到的BUG

目录

• 1. zip argument #1 must support iteration
• 2. torch.nn.DataParallel
• 3. model.state_dict()

1. zip argument #1 must support iteration

在多gpu训练的时候，自动把你的batch_size分成n_gpu份，每个gpu跑一些数据， 最后再合起来。我之所以出现这个bug是因为返回的时候 返回了一个常量。。

2. torch.nn.DataParallel

在使用torch.nn.DataParallel时候，要先把模型放在gpu上，再进行parallel。

3. model.state_dict()

一般在现有的网络加载预训练模型通常是找到预训练模型在现有的model里面的参数，然后model进行更新，遇到一个bug， 发现加载预训练模型的时候， 效果很差，跟参数没有更新一样，找了一大顿bug，最后才发现，之前是单gpu进行的预训练，现在的模型使用的是多gpu， 打印现在模型的参数你会发现他所有的参数前面都加了一个module. 所以向以前一样更新，没有一个参数会被更新，因此写了一个万能模型参数加载函数。

pretrained_dict = checkpoint['state_dict']
model_dict = self.model.state_dict()
if checkpoint['config']['n_gpu'] > 1 and self.config['n_gpu'] == 1:
    new_dict = OrderedDict()
    for k, v in pretrained_dict.items():
        name = k[7:]
        new_dict[name] = v
    pretrained_dict = new_dict
elif checkpoint['config']['n_gpu'] == 1 and self.config['n_gpu'] > 1:
    new_dict = OrderedDict()
    for k, v in pretrained_dict.items():
        name = "module."+k
        new_dict[name] = v
    pretrained_dict = new_dict
print("The pretrained model's para is following")
for k, v in pretrained_dict.items():
    print(k)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
self.model.load_state_dict(model_dict)