0

将 Python 3.6 与 Pytorch 1.3.1 一起使用。我注意到当整个模块被导入另一个模块时，一些保存的 nn.Modules 无法加载。举个例子，这里是一个最小工作示例的模板。

#!/usr/bin/env python3
#encoding:utf-8
# file 'dnn_predict.py'

from torch import nn
class NN(nn.Module):##NN network
    # Initialisation and other class methods

networks=[torch.load(f=os.path.join(resource_directory, 'nn-classify-cpu_{fold}.pkl'.format(fold=fold))) for fold in range(5)]
...
if __name__=='__main__':
    # Some testing snippets
    pass

当我直接在 shell 中运行它时，整个文件工作得很好。但是，当我想使用该类并使用此代码将神经网络加载到另一个文件中时，它会失败。

#!/usr/bin/env python3
#encoding:utf-8
from dnn_predict import *

错误读取AttributeError: Can't get attribute 'NN' on <module '__main__'>

在 Pytorch 中加载保存的变量或导入模块是否与其他常见的 Python 库不同？一些帮助或指向根本原因的指针将不胜感激。

4

1 回答 1

1

当您保存模型时torch.save(model, PATH)，整个对象会被序列化pickle，这不会保存类本身，而是包含该类的文件的路径，因此在加载模型时需要完全相同的目录和文件结构才能找到正确的类. 运行 Python 脚本时，该文件的模块是__main__，因此如果要加载该模块，则NN必须在正在运行的脚本中定义类。

那是非常不灵活的，所以推荐的做法是不保存整个模型，而是只保存状态字典，它只保存模型的参数。

# Save the state dictionary of the model
torch.save(model.state_dict(), PATH)

之后，可以加载状态字典并将其应用于您的模型。

from dnn_predict import NN

# Create the model (will have randomly initialised parameters)
model = NN()

# Load the previously saved state dictionary
state_dict = torch.load(PATH)

# Apply the state dictionary to the model
model.load_state_dict(state_dict)

有关状态字典和保存/加载模型的更多详细信息：PyTorch - 保存和加载模型

于 2020-05-04T13:04:13.350 回答