Python 保存PyTorch模型，但不访问模型类代码

如果不需要在某处定义模型类，如何保存PyTorch模型

---

免责声明：

在中，没有用于保存模型而不访问模型类代码的解决方案（或工作解决方案）

在没有类访问权限的情况下，没有用于保存模型的解决方案（或工作解决方案）

你想存什么就存什么

您可以保存模型，

torch.save（模型，文件路径）

。它保存模型对象本身

您可以只保存模型状态dict

torch.save（model.state\u dict（），filepath）

此外，您可以保存任何您喜欢的内容，因为

torch.save

只是基于pickle的保存

state = {
    'hello_text': 'just the optimizer sd will be saved',
    'optimizer': optimizer.state_dict(),

}
torch.save(state, filepath)

---

你可以过一会儿再查

我建议您将pytorch模型转换为onnx并保存它。这可能是在不访问类的情况下存储模型的最佳方式。

提供了一个PyTorch核心开发人员（smth）的正式答案：

在没有代码的情况下加载pytorch模型有一些限制

第一个限制： 我们只保存类定义的源代码。我们不保存超出此范围的内容（如类所引用的包源代码）

例如：

import foo

class MyModel(...):
    def forward(input):
        foo.bar(input)

此处，包

foo

未保存在模型检查点中

第二限制： 健壮地序列化python构造有一些限制。例如，默认pickler无法序列化lambda。有一些帮助程序包可以序列化比标准更多的python构造，但它们仍然有局限性。Dill 25就是这样一个包

---

鉴于这些限制，没有Trr.Load工作的鲁棒方法。没有原始源文件。

< P>如果你计划用PyTrar库（即Python、C++或其他平台支持的PyTrk）进行推理，那么最好的方法是通过.< /P> 我认为最简单的方法是使用

trace=torch.jit.trace（model，典型输入）

，然后使用

torch.jit.save（trace，path）

。然后可以使用

torch.jit.load（path）

加载跟踪模型

这里有一个非常简单的例子。我们制作了两个文件：

train.py

：

导入火炬
类别模型（torch.nn.Module）：
定义初始化（自）：
超级（模型，自我）。\uuuu初始化
self.linear=torch.nn.linear（4,4）
def前进（自身，x）：
x=火炬释放（自线性（x））
返回x
模型=模型（）
x=火炬.浮子张量（[[0.2,0.3,0.2,0.7]，[0.4,0.2,0.8,0.9]）
使用手电筒。无梯度（）
印刷品（型号（x））
跟踪单元=火炬准时跟踪（型号，（x））
torch.jit.save（跟踪单元格，“model.pth”）

推断.py

：

导入火炬
x=火炬.浮子张量（[[0.2,0.3,0.2,0.7]，[0.4,0.2,0.8,0.9]）
已加载的\u trace=torch.jit.load（“model.pth”）
使用手电筒。无梯度（）
打印（加载的跟踪（x））

按顺序运行这些命令将得到以下结果：

python train.py
张量（[[0.0000,0.1845,0.2910,0.2497]，
[0.0000, 0.5272, 0.3481, 0.1743]])
python推断.py
张量（[[0.0000,0.1845,0.2910,0.2497]，
[0.0000, 0.5272, 0.3481, 0.1743]])

结果是一样的，所以我们很好。（注意，由于nn.线性层初始化的随机性，每次的结果都会不同）

---

TorchScript提供了更复杂的体系结构和图形定义（包括if语句、while循环等），可以保存在单个文件中，而无需在推理时重新定义图形。有关更高级的可能性，请参阅文档（上面链接）。

我希望存在一些解决方法，因为tensorflow中有这样的选项。我对onnx不感兴趣，因为它有一些限制…并且不支持一些pytorch功能…使用torch脚本的缺点是什么？主要的缺点是您仍然需要某种pytorch环境。另外，如果你想继续训练跟踪，我想那将是非常困难/不可能的。有时也可能有点问题/难以调试。但这基本上是pytorch的答案，即在tensorflow中轻松保存整个图形。每一个版本都在改进，在imo中已经非常好了。