torchScript模型 和 torch 通过模型代码创建模型区别，torchScript模型转化方法和加载方法

目录

TorchScript模型与Torch模型代码创建的区别

1. 定义

TorchScript模型

通过Torch模型代码创建的模型

2. 使用场景和目的

TorchScript模型

通过Torch模型代码创建的模型

3. 灵活性和调试

TorchScript模型

通过Torch模型代码创建的模型

4. 兼容性和维护

TorchScript模型

通过Torch模型代码创建的模型

转换PyTorch模型为TorchScript模型的方法

1. 使用tracing（跟踪）

2. 使用scripting（脚本化）

3. 混合前端（Combining Tracing and Scripting）

加载TorchScript模型的方法

1. 使用torch.jit.load函数

2. 在不同的设备上加载模型

3. 加载到指定的作用域

---

TorchScript模型与Torch模型代码创建的区别

1. 定义

TorchScript模型

• TorchScript是PyTorch的一个子集，可以通过两种方式创建：跟踪（tracing）和脚本（scripting）。
• TorchScript提供了一种将PyTorch模型序列化的方法，允许它们在不依赖Python解释器的环境中运行，例如在C++程序中。
• TorchScript模型可以在Torch JIT（Just-In-Time）编译器中运行，这有助于优化模型的执行速度和内存使用。

通过Torch模型代码创建的模型

• 这是指直接使用PyTorch框架通过Python代码定义的模型，通常使用torch.nn.Module类及其子类来构建模型架构。
• 这样的模型在Python环境中动态执行，依赖于Python解释器。

2. 使用场景和目的

TorchScript模型

• 用于模型部署：当需要将模型导出到生产环境，特别是非Python环境时，使用TorchScript是很有用的。
• 跨平台运行：TorchScript模型可以在不同的平台和设备上运行，不受Python环境限制。
• 性能优化：TorchScript模型可以通过JIT编译器进行优化，提高运行效率。

通过Torch模型代码创建的模型

• 用于模型开发和训练：在开发和训练阶段，模型通常直接用Python代码定义，因为这样更灵活，易于调试。
• 交互式开发：Python环境支持交互式开发，可以即时测试和修改模型。

3. 灵活性和调试

TorchScript模型

• 灵活性较低：转换为TorchScript的模型可能需要去掉依赖于Python的某些动态特性，以保证模型可以被序列化。
• 调试困难：TorchScript模型不易于调试，因为它们是在Python环境之外运行的。

通过Torch模型代码创建的模型

• 灵活性高：可以使用Python的全部功能，包括动态图构建和各种控制流。
• 易于调试：在Python环境中可以使用标准的调试工具，如pdb或IDE内置的调试器。

4. 兼容性和维护

TorchScript模型

• 兼容性好：TorchScript模型可以在不同版本的PyTorch和不同的系统中运行，有助于长期维护。
• 维护成本：一旦模型被转换为TorchScript，对模型的进一步修改可能需要重新转换。

通过Torch模型代码创建的模型

• 版本依赖：模型可能依赖于特定版本的PyTorch和第三方库。
• 维护灵活：可以直接修改Python代码来更新或维护模型。

总结来说，TorchScript模型适合于模型的优化、部署和跨平台运行，而直接通过PyTorch代码创建的模型则更适合于模型的开发和训练阶段。选择哪种方式取决于具体的应用场景和需求。

转换PyTorch模型为TorchScript模型的方法

1. 使用tracing（跟踪）

• 跟踪是一种通过运行模型的正向传播来记录操作的方法，这通常用于没有控制流（如if语句和循环）的模型。
• 使用torch.jit.trace函数，你可以传入模型（nn.Module对象）和一组代表输入的示例张量。
• 跟踪过程会执行一次模型的正向传播，并记录所有的操作。
• 返回的是一个ScriptModule，它是一个TorchScript模型，可以独立于原始Python代码运行。

import torch

# 假设我们有一个已经训练好的模型
model = MyModel()

# 准备一个输入张量example_input
example_input = torch.rand(1, 3, 224, 224)

# 使用tracing将模型转换为TorchScript
traced_script_module = torch.jit.trace(model, example_input)

# 保存TorchScript模型供以后使用或部署
traced_script_module.save("model.pt")
 

2. 使用scripting（脚本化）

• 脚本化是一种将PyTorch模型及其控制流转换为TorchScript的方法，它通过分析Python代码来创建一个静态图。
• 使用torch.jit.script函数可以将一个nn.Module对象转换为ScriptModule。
• 脚本化不需要输入张量，因为它分析的是代码而非执行过程。
• 这种方法适用于模型中包含复杂控制流（如if语句、循环和递归函数）的情况。

import torch

# 假设我们有一个已经训练好的模型
model = MyModel()

# 使用scripting将模型转换为TorchScript
script_module = torch.jit.script(model)

# 保存TorchScript模型供以后使用或部署
script_module.save("model.pt")

3. 混合前端（Combining Tracing and Scripting）

• 对于一些模型，可能需要结合跟踪和脚本化的方法来转换，这称为混合前端。
• 在这种情况下，可以将模型的某些部分标记为脚本化（使用torch.jit.script装饰器），而其他部分则通过跟踪转换。
• 这允许在不牺牲控制流的情况下，对模型的特定部分进行优化。

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # ... 初始化 ...

    @torch.jit.script_method
    def forward(self, x):
        # ... 实现含有控制流的前向传播 ...
        return x

# 创建模型实例
model = MyModel()

# 使用tracing转换模型的其他部分
example_input = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example_input)

# 保存TorchScript模型
traced_script_module.save("model.pt")

转换模型为TorchScript格式后，可以通过调用.save()方法将其保存为一个文件，这个文件可以在不同的环境中加载和运行，无需Python解释器。

加载TorchScript模型的方法

1. 使用torch.jit.load函数

• torch.jit.load是一个用于加载TorchScript模型的函数，它接受一个指向序列化模型文件的路径。
• 加载后，返回一个ScriptModule对象，该对象可以像常规的PyTorch模型一样使用。

import torch

# 加载先前保存的TorchScript模型
model = torch.jit.load("model.pt")

# 使用加载的模型进行推理
example_input = torch.rand(1, 3, 224, 224)
output = model(example_input)
 

2. 在不同的设备上加载模型

• 在加载模型时，可以指定模型运行的设备，例如CPU或GPU。
• 使用map_location参数来指定加载模型时张量的设备位置。

# 加载模型到CPU
model = torch.jit.load("model.pt", map_location=torch.device('cpu'))

# 或者加载模型到指定的GPU设备
model = torch.jit.load("model.pt", map_location=torch.device('cuda:0'))

3. 加载到指定的作用域

• 如果需要在特定的作用域中加载模型，比如一个函数或类的内部，可以使用torch.jit.load的_extra_files参数加载额外的文件。

# 加载模型和附加文件
extra_files = {'extra_file.txt': 'r'}
model = torch.jit.load("model.pt", _extra_files=extra_files)

加载TorchScript模型后，可以直接使用该模型执行前向传播，进行推理或其他操作。如果模型是在GPU上训练的，确保在相同或兼容的设备上加载模型，以避免设备不匹配的问题。如果需要在不同的设备之间迁移模型，使用map_location参数来指定目标设备。