python - Pytorch：使用 DataParallel 进行训练并在 cpu 上进行测试

Question

我有一个使用我在 pytorch 中的其他预训练模型之一的训练模型。对于测试，我想在 cpu 而不是 gpu 上进行测试。问题是，由于我在训练过程中使用了 Dataparallel 函数，因此将其转换为 cpu 对我来说并不简单。

我的训练模型代码如下：

class DGCNN_tl(nn.Module):
    def __init__(self):
        super(DGCNN_tl, self).__init__()
        self.model = DGCNN()
        if torch.cuda.device_count() > 1:
            print(str(torch.cuda.device_count())+" GPUs detected!")
            self.model = nn.DataParallel(self.model)
            patch_replication_callback(self.model) #sync_batch_norm
            self.model.to(device)
        self.model.load_state_dict(torch.load("teeth_seg2_pretrained.pt"))
        
        for param in self.model.parameters(): #making all layers fine tuneable
            param.requires_grad = True
            
    def forward(self, inputs, l):
        outputs,t,feature = self.model(inputs, l)
        return outputs, t, feature

以下是我尝试将其转换为 cpu 进行测试的方法：

dgcnn_lower = DGCNN_tl()
# original saved file with DataParallel
model_old = torch.load('teeth_seg2_lower_0.953iou.pt', map_location=torch.device('cpu'))

# create new OrderedDict that does not contain `module.`
from collections import OrderedDict

new_state_dict = OrderedDict()
for k, v in model_old.items():
    name = k[7:] # remove `module.`
    new_state_dict[name] = v

# load params
dgcnn_lower.load_state_dict(new_state_dict)

这给出了一个错误说

    Missing key(s) in state_dict: "model.module.transform_net.bn1.weight", "model.module.transform_net.bn1.bias"

“teeth_seg2_pretrained.pt”：这是我的预训练模型，用于训练我当前的模型 “teeth_seg2_lower_0.953iou.pt”：这是我想在 cpu 上测试的当前模型

有人可以给我一个建议或帮助来解决这个问题吗？谢谢！！

标签： pythondeep-learningpytorch