问题描述

训练模型时，分阶段训练，第二阶段加载第一阶段训练好的模型的参数，接着训练
第一阶段训练，含有代码

    if (train_on_gpu):if torch.cuda.device_count() > 1:net = nn.DataParallel(net)net = net.to(device)

第二阶段训练，含有代码

    if (train_on_gpu):if torch.cuda.device_count() > 1:netT = nn.DataParallel(netT)netS = nn.DataParallel(netS)netT = netT.to(device)netS = netS.to(device)
-----orig_state_dict = torch.load("../models/model.pth")['net']new_state_dict = OrderedDict()for k, v in orig_state_dict.items():name = k.replace('module.', '')new_state_dict[name] = vnetT.load_state_dict(new_state_dict)
-----    for param in netS.parameters():param.requires_grad = True#对源模型 netS 的 pose_fc1 层所有参数都设置为不需要进行反向传播更新。for param in netS.pose_fc1.parameters():param.requires_grad = False

结果报错

RuntimeError: Error(s) in loading state_dict for DataParallel: Missing key(s) in state_dict: "module.feat.block1.layer.0.conv1.weight", ...........

解决方案：

注释掉nn.DataParallel(）

    if (train_on_gpu):if torch.cuda.device_count() > 1:pass#netT = nn.DataParallel(netT)#netS = nn.DataParallel(netS)netT = netT.to(device)netS = netS.to(device)
-----orig_state_dict = torch.load("../models/model.pth")['net']new_state_dict = OrderedDict()for k, v in orig_state_dict.items():name = k.replace('module.', '')new_state_dict[name] = vnetT.load_state_dict(new_state_dict)
-----    for param in netS.parameters():param.requires_grad = True#对源模型 netS 的 pose_fc1 层所有参数都设置为不需要进行反向传播更新。for param in netS.pose_fc1.parameters():param.requires_grad = False

原因分析

可能是模型在第一阶段和第二阶段训练设置不一致导致的问题，比如第一阶段用双卡训练，第二阶段用单卡训练
这时不能第一阶段和第二阶段都用nn.DataParallel()

我还不太清楚，我猜的

而net加载的网络结构没有"module"，而第一阶段保存的模型因为使用了net = nn.DataParallel(net)，保存的参数的key有"module"，那么在第二阶段加载的时候就要把key前面的"module"去掉，才能正确加载模型参数到网络里面，继续训练