### 将Python脚本中的模型保存为PyTorch的pt格式 在PyTorch中，模型可以通过两种方式保存为`.pt`文件：一种是保存整个模型（包括结构和权重），另一种是仅保存模型的权重。以下分别介绍这两种方法。 #### 1. 保存整个模型 当需要保存整个模型时，可以使用`torch.save()`函数将模型对象直接保存到文件中。这种方法会保存模型的结构和权重，适合于需要重新加载模型并直接进行推理或训练的情况。 ```python import torch from test import TheModelClass # 初始化模型 model = TheModelClass() # 保存整个模型为 .pt 文件 torch.save(model, "models/entire_model.pt") ``` 通过这种方式保存的模型可以通过以下代码加载[^1]： ```python # 加载整个模型 model = torch.load("models/entire_model.pt") model.eval() # 设置模型为评估模式 ``` #### 2. 仅保存模型的权重 如果只需要保存模型的权重（推荐的方式），可以使用`torch.save()`保存模型的`state_dict`。这种方式更加高效且灵活，因为模型的结构可以在加载时重新定义。 ```python import torch from test import TheModelClass # 初始化模型 model = TheModelClass() # 保存模型的 state_dict 为 .pt 文件 torch.save(model.state_dict(), "models/model_state_dict.pt") ``` 加载模型权重时，需要先定义模型的结构，然后加载保存的权重[^1]： ```python # 定义模型结构 model = TheModelClass() # 加载权重 model.load_state_dict(torch.load("models/model_state_dict.pt")) model.eval() # 设置模型为评估模式 ``` #### 3. 动态图与静态图的区别 需要注意的是，上述方法适用于动态图模型。如果希望将模型转换为静态图（即TorchScript），可以使用`torch.jit.script()`或`torch.jit.trace()`生成`.pt`文件[^2]。例如： ```python import torch from test import TheModelClass # 初始化模型 model = TheModelClass() model.eval() # 使用 trace 方法生成静态图模型 example_input = torch.rand(1, 3, 224, 224) # 示例输入 traced_script_module = torch.jit.trace(model, example_input) # 保存为 .pt 文件 traced_script_module.save("models/traced_model.pt") ``` 加载静态图模型时，可以直接使用`torch.jit.load()`： ```python import torch # 加载静态图模型 model = torch.jit.load("models/traced_model.pt") model.eval() ``` ### 注意事项 - 如果保存的是整个模型，则加载时需要确保环境中有模型的定义文件。 - 如果保存的是`state_dict`，则需要手动定义模型结构后加载权重。 - 对于C++应用，推荐使用TorchScript保存的`.pt`文件，因为它可以直接被LibTorch加载和执行。