torch.jit.trace 和 torch.jit.script 如何使用

torch.jit.trace和torch.jit.script是PyTorch的两种模型序列化工具，用于将PyTorch模型序列化为可保存和加载的文件格式。它们的使用方法如下：

1. torch.jit.trace

torch.jit.trace用于将PyTorch模型转换为TorchScript，可以使得模型在C++中运行。使用方法如下：

import torch

# 定义模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.relu = torch.nn.ReLU()
        self.conv2 = torch.nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.fc = torch.nn.Linear(32 * 32 * 32, 10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = x.view(-1, 32 * 32 * 32)
        x = self.fc(x)
        return x

# 实例化模型
model = MyModel()

# 定义输入数据
input_data = torch.rand(1, 3, 64, 64)

# 将模型转换为TorchScript
traced_model = torch.jit.trace(model, input_data)

# 保存TorchScript模型
traced_model.save('my_model.pt')

2. torch.jit.script

torch.jit.script用于将PyTorch模型转换为TorchScript，可以使得模型在C++中运行。使用方法如下：

import torch

# 定义模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.relu = torch.nn.ReLU()
        self.conv2 = torch.nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.fc = torch.nn.Linear(32 * 32 * 32, 10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = x.view(-1, 32 * 32 * 32)
        x = self.fc(x)
        return x

# 实例化模型
model = MyModel()

# 定义输入数据
input_data = torch.rand(1, 3, 64, 64)

# 将模型转换为TorchScript
scripted_model = torch.jit.script(model)

# 保存TorchScript模型
scripted_model.save('my_model.pt')

以上是torch.jit.trace和torch.jit.script的使用方法。需要注意的是，如果模型中使用了一些Python特性或库，如if语句、for循环等，则只能使用torch.jit.script进行转换，而不能使用torch.jit.trace。