import torch import torch.nn as nn # 定义一个线性模型 class LinearModel(nn.Module): def __init__(self, input_size, output_size): super(LinearModel, self).__init__() self.linear = nn.Linear(input_size, output_size) def forward(self, x): return self.linear(x) # 创建模型实例 model = LinearModel(input_size=10, output_size=1) print("Model:", model) # 方法 1：保存和加载 state_dict # 保存模型的 state_dict torch.save(model.state_dict(), './linear_model_state_dict.pth') # 创建一个新的模型实例 new_model = LinearModel(input_size=10, output_size=1) # 加载 state_dict 到新模型 new_model.load_state_dict(torch.load('./linear_model_state_dict.pth')) # 方法 2：保存和加载整个模型 # 保存整个模型 torch.save(model, './linear_model.pth') # 加载整个模型 loaded_model = torch.load('./linear_model.pth')