无法在 Pytorch Lightning 中加载自定义的预训练权重答案

【问题标题】：Unable to load custom pretrained weight in Pytorch Lightning无法在 Pytorch Lightning 中加载自定义的预训练权重
【发布时间】：2022-01-15 18:56:34
【问题描述】：

我想用我的小数据集重新训练自定义模型。我可以加载预训练的权重 (.pth) 并在 Pytorch 中运行它。但是，我需要更多功能并将代码重构为 Pytorch Lightning，但我不知道如何将预训练的权重加载到 Pytorch Lightning 模型中。

请在下面查看我的代码的详细信息：

class BDRAR(nn.Module):
    def __init__(self):
        super(BDRAR, self).__init__()
        resnext = ResNeXt101()
        self.layer0 = resnext.layer0
        self.layer1 = resnext.layer1
        self.layer2 = resnext.layer2
        self.layer3 = resnext.layer3
        self.layer4 = resnext.layer4

Pytorch 闪电代码：

class liteBDRAR(pl.LightningModule):
    def __init__(self):
        super(liteBDRAR, self).__init__()
        self.model = BDRAR()
        print('Model Created!')

    def forward(self, x):
        return self.model(x)

Pytorch 闪电运行：

    path = './ckpt/BDRAR/3000.pth'
    bdrar = liteBDRAR.load_from_checkpoint(path,  strict=False)
    trainer = pl.Trainer(fast_dev_run=True, gpus=1)
    trainer.fit(bdrar)

错误：

keys = model.load_state_dict(checkpoint["state_dict"], strict=strict)
**KeyError: 'state_dict'**

我将不胜感激。

谢谢。

【问题讨论】：

标签： deep-learning pytorch pytorch-lightning

【解决方案1】：

您收到此错误的原因是您试图将 PyTorch 的模型权重加载到 Lightning 模块中。使用 Lightning 保存检查点时，您不仅可以保存模型状态，还可以保存大量其他信息（请参阅 here）。

您正在寻找的是以下内容：

path = './ckpt/BDRAR/3000.pth'
bdrar = liteBDRAR()
bdrar.model.load_state_dict(torch.load(path))

【讨论】：

非常感谢！我会试试这个，然后告诉你结果。
我收到此错误：RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.
这个问题不再与此相关。我建议你看看here
谢谢@Jules！

【解决方案2】：

您的.pth 文件可能已经是state_dict。尝试在您的闪电课程中加载预训练的重量。

class liteBDRAR(pl.LightningModule):
    def __init__(self):
        super(liteBDRAR, self).__init__()
        self.model = BDRAR()
        print('Model Created!')

    def load_model(self, path):
        self.model.load_state_dict(torch.load(path, map_location='cuda:0'), strict=False)

path = './ckpt/BDRAR/3000.pth'
model = liteBDRAR()
model.load_model(path)

【讨论】：

这是我尝试时得到的错误：self.model = torch.load_from_dict(torch.load(path), strict=False) AttributeError: module 'torch' has no attribute 'load_from_dict'
对不起，我更新了功能。应该是load_state_dict。
这是我得到的错误：AttributeError: module 'torch' has no attribute 'load_state_dict'
又是我的错 :( 改用self.model.load_state_dict(torch.load(path), strict=False)。更新了答案。
谢谢。我试过了，这次得到了一个不同的错误：RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.

【解决方案3】：

那些预训练的权重属于class BDRAR(nn.Module)。也就是说，你的闪电模块的model 参数中的类。

LightningModule liteBDRAR() 充当 Pytorch 模型的包装器（位于 self.model）。您需要将权重加载到您的闪电模块内的 pytorch 模型上。正如@Jules 和@Dharman 提到的，您需要的是：

path = './ckpt/BDRAR/3000.pth'
bdrar = liteBDRAR()
bdrar.model.load_state_dict(torch.load(path))

【讨论】：

谢谢@Sarita。当我尝试 Jules 和 Dharman 的建议时，我遇到了运行时错误。 joe32140 的最后一个建议奏效了。