加载 PyTorch NN 模型的检查点时出现异常

Question

调用以下单元格中定义的函数时，抛出异常'TypeError: forward() 接受 2 个位置参数，但给出了 9 个' 这个document 提供了更多细节

def load_checkpoint(chkptJP):
checkpoint = torch.load(chkptJP)
model2 = model1(checkpoint['input_size'],
              checkpoint['output_size'],
              checkpoint['fc1'],
              checkpoint['fc2'],
              checkpoint['optimizer_state_dict'],
              checkpoint['epoch'],
              checkpoint['class_to_idx'],
              checkpoint['learning_rate'])
model2.load_state_dict(checkpoint['state_dict'])
return model2

写出检查点的代码如下：

checkpoint ={'input_size':512,
         'output_size':102,
         'fc1':256,
         'fc2':102,
         'state_dict': model.state_dict(),
         'optimizer_state_dict': optimizer.state_dict(),
         'epoch': epoch+1,
         'class_to_idx': model.class_to_idx,
         'learning_rate': 0.003}
torch.save(checkpoint,chkptJP)

Answer 1

您的错误表明model1 是一个已经实例化的网络，而它应该是一个类。有关全面信息，请参阅official documentation about saving（如有疑问，请始终参考）。我将在整个答案中链接到它，因此请务必查看并了解发生了什么。

保存常规检查点

您的代码保存了general checkpoint。您可以通过这种方式保存任何字典以及您想要的任何信息（基本上是Python's pickle，您也可以类似地对其进行调整）。您的信息很多，其中一些与模型本身无关。

加载常规检查点

正如您所做的那样，您可以通过torch.load 加载所有这些数据。由于您保存了state_dict（权重），而不是整个Model（代码的外观），因此您必须使用随机权重创建模型并在之后加载它们。

这段代码应该没问题：

def load_checkpoint(chkptJP):
    checkpoint = torch.load(chkptJP)
    model = ModelClass(
        checkpoint["input_size"],
        checkpoint["output_size"],
        checkpoint["fc1"],
        checkpoint["fc2"],
        checkpoint["optimizer_state_dict"],
        checkpoint["epoch"],
        checkpoint["class_to_idx"],
        checkpoint["learning_rate"],
    )
    model.load_state_dict(checkpoint["state_dict"])
    return model

注意 ModelClass 必须是类，而不是您在此处所做的对象。如果model1 是一个对象，运行model1(arg1, ..., arg9) 将调用它的__call__ 方法，如果model1 是torch.nn.Module 的一个实例，则该方法又是一个包装的forward 方法。 ModelClass 在您的代码中应该是这样的（并且可能在某处定义）：

import torch


class ModelClass(torch.nn.Module):
    def __init__(
        self,
        input_size,
        output_size,
        fc1,
        fc2,
        optimizer_state_dict,
        epoch,
        class_to_idx,
        learning_rate,
    ):
        # Your initialization code here
        ...

    def forward(tensor):
        # Your forward pass here
        ...

如果您在任何地方都没有 ModelClass，则必须单独保存整个模型（例如，torch.save(model) 而不是 torch.save(model.state_dict())) 并将其作为一个整体加载（torch.load(PATH) 而不是 chkp=torch.load(PATH)，然后是 model.load_state_dict在实例上调用）