pytorch - 如何在 PyTorch 中保存模型架构？

Question

我知道我可以通过 torch.save(model.state_dict(), FILE) 保存模型或 torch.save(model, FILE) .但是它们都没有保存模型的架构。

那么我们如何在 PyTorch 中保存模型的架构，比如创建一个 .pb Tensorflow 中的文件？我想对我的模型应用不同的调整。如果我不能保存模型的架构，我有什么比每次复制整个类定义并创建一个新类更好的方法吗？

Answer 1

您可以引用this文章了解如何保存分类器。要对模型进行调整，您可以做的是创建一个新模型，它是现有模型的子代。

class newModel( oldModelClass):
    def __init__(self):
        super(newModel, self).__init__()

通过这种设置，newModel 拥有所有层以及 oldModelClass 的前向函数。 .如果您需要进行调整，您可以在 __init__ 中定义新图层。函数，然后编写一个新的前向函数来定义它。