PyTorch模型保存技巧

pytorch模型保存技巧
在PyTorch中，保存模型主要有两种方法：保存模型的结构和权重。保存模型的结构可以帮助你在之后加载模型时，可以知道模型的结构和参数的位置，而保存模型的权重则可以使你在之后加载模型时，能够立即获得之前训练好的权重。在大多数情况下，你可能需要同时保存模型的结构和权重。

保存模型结构

保存模型结构非常简单，你只需要将模型的状态_dict()保存为文件即可。但是，如果你想要在之后加载模型时知道模型的结构，你需要保存一些额外的信息，如模型的类名和层名称。你可以使用torch.jit.script()或者torch.jit.trace()将模型转换为TorchScript，然后保存为文件。
例如：

import torch
import torchvision.models as models
# 创建一个预训练的resnet模型
model = models.resnet50(pretrained=True)
# 将模型转换为TorchScript
traced_script_module = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
# 保存模型
traced_script_module.save("resnet.pt")

在上面的例子中，我们首先创建了一个预训练的ResNet模型，然后将其转换为TorchScript，最后将其保存为文件。如果你之后想要加载这个模型，你只需要使用torch.jit.load()函数即可。

保存模型权重

保存模型权重也非常简单。你只需要将模型的state_dict()保存为文件即可。例如：

# 保存模型权重
torch.save(model.state_dict(), "model_weights.pt")

在上面的例子中，我们使用torch.save()函数将模型的权重保存为文件。如果你之后想要加载这个模型的权重，你只需要使用model.load_state_dict()函数即可。例如：

# 加载模型权重
model.load_state_dict(torch.load("model_weights.pt"))

在上面的例子中，我们使用torch.load()函数加载之前保存的权重，然后将其赋值给模型的state_dict()。注意，你需要确保在加载权重之前，你已经创建了模型的架构。如果你还没有创建模型的架构，你需要先创建它，然后才能加载权重。