pytorch保存和导入模型

Pytorch官方的加载和保存模型的方式有两种:

1、保存和加载整个模型

这种方式再重新加载的时候不需要自定义网络结构,保存时已经把网络结构保存了下来,比较死板不能调整网络结构。

注:torch.load 返回的是一个 OrderedDict

torch.save(model_object, 'model.pkl')
model = torch.load('model.pkl') 

2、仅保存和加载模型参数(推荐使用) state_dict

这种方式再重新加载的时候需要自己定义网络,并且其中的参数名称与结构要与保存的模型中的一致(可以是部分网络,比如只使用VGG的前几层),相对灵活,便于对网络进行修改。
 

torch.save(model_object.state_dict(), 'params.pkl')
model_object.load_state_dict(torch.load('params.pkl'))

特殊情况:

一、当保存的模型和导入的模型不在一个位置时

假设我们只保存了模型的参数(model.state_dict())到文件名为modelparameters.pth, model = Net()

1. cpu -> cpu或者gpu -> gpu:

checkpoint = torch.load('modelparameters.pth')

model.load_state_dict(checkpoint)

2. cpu -> gpu 1

torch.load('modelparameters.pth', map_location=lambda storage, loc: storage.cuda(1))

3. gpu 1 -> gpu 0

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'})

4. gpu -> cpu

torch.load('modelparameters.pth', map_location=lambda storage, loc: storage)

二、多GPU训练

例如我们创建了一个多GPU训练模型:

model = MyVggNet()
#多GPU并行
model = nn.DataParallel(model).cuda()

此时相当于在原来的模型外面加了一层支持GPU运行的外壳,真正的模型对象为:real_model = model.module。

所以在保存模型的时候注意,如果保存的时候是否带有这层加的外壳,加载的时候也是带有的,如果保存的是真实的模型,加载的也是真是的模型。因为加了module壳的模型在CPU上是不能运行的,因此建议保存真实模型。

1. 第一种方式

模型保存(保存带外壳模型的参数):

real_model = model.module
torch.save(real_model.state_dict(),os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_weight.pth"))

cpu上加载使用:

args.weight=checkpoint/cos_mnist_10_weight.pth
map_location = lambda storage, loc: storage
model.load_state_dict(torch.load(args.weight,map_location=map_location))


2. 第二种方式

模型保存(保存真正模型的参数):

real_model = model.module
save_model(real_model, os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_weight_cpu.pth"))
# 自定义的函数
def save_model(model,filename):
    state = model.state_dict()
    for key in state: state[key] = state[key].clone().cpu()
    torch.save(state, filename)

cpu上加载使用:

args.weight=checkpoint/cos_mnist_10_weight_cpu.pth
model.load_state_dict(torch.load(args.weight))


3. 第三种方式


模型保存(保存整个带壳的模型)

real_model = model.module
torch.save(real_model, os.path.join(args.save_path,"cos_mnist_"+str(epoch+1)+"_whole.pth"))

cpu上加载使用:

args.weight=checkpoint/cos_mnist_10_whole.pth
map_location = lambda storage, loc: storage
model = torch.load(args.weight,map_location=map_location)

 

  • 0
    点赞
  • 16
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
### 回答1: 要在 PyTorch保存训练好的模型,可以使用 PyTorch 的内置函数 `torch.save()`。下面是保存模型的基本步骤: 1. 在训练完成后,将模型保存为字典类型,包括模型的状态字典和其他必要的参数。 ```python model_state = { 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, 'loss': loss } ``` 2. 使用 `torch.save()` 函数将模型保存到指定的文件路径。 ```python torch.save(model_state, 'model.pt') ``` 3. 若要加载模型,请使用 `torch.load()` 函数来加载状态字典并将其赋值给模型。 ```python model_state = torch.load('model.pt') model.load_state_dict(model_state['state_dict']) ``` 这些步骤可以适用于各种不同的 PyTorch 模型。另外,如果您想要更好地控制保存的内容和格式,可以使用 PyTorch 的序列化函数将模型保存为自定义格式的文件。 ### 回答2: 要保存训练好的模型,我们可以使用PyTorch的torch.save()函数。该函数接受两个参数:第一个是要保存模型的状态字典,第二个是保存路径。 首先,我们需要定义一个模型的实例,然后在训练过程中,使用optimizer进行模型的优化。当训练完成后,我们可以使用torch.save()函数将模型的状态字典保存下来。 下面是保存模型的示例代码: ```python import torch import torch.nn as nn import torch.optim as optim import torchvision.models as models # 定义一个模型实例 model = models.resnet18() # 定义优化器 optimizer = optim.Adam(model.parameters(), lr=0.001) # ... 在训练过程中进行模型训练 ... # 训练完成后,保存模型的状态字典 torch.save(model.state_dict(), 'model.pth') ``` 以上代码中,我们使用了ResNet18作为示例模型,Adam作为优化器。在训练完成后,使用torch.save()函数将模型的状态字典保存到'model.pth'文件中。 注意,只保存模型的状态字典是因为它包含了模型的参数和缓存信息,但不包含模型的结构。在加载模型时,我们需要先创建一个模型实例,然后使用torch.load()函数加载保存模型状态字典。接下来,通过调用模型实例的load_state_dict()方法加载模型的状态字典,使模型获取保存的参数与缓存信息。 希望以上回答对您有所帮助! ### 回答3: 在PyTorch保存训练好的模型可以通过以下步骤完成: 1. 导入必要的库: ``` import torch import torch.nn as nn import torch.optim as optim ``` 2. 定义模型结构: ``` class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.fc = nn.Linear(input_size, output_size) def forward(self, x): x = self.fc(x) return x ``` 3. 创建模型实例: ``` model = MyModel() ``` 4. 定义损失函数和优化器: ``` criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) ``` 5. 训练模型: ``` for epoch in range(num_epochs): # 前向传播 outputs = model(inputs) loss = criterion(outputs, labels) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() ``` 6. 保存模型: ``` torch.save(model.state_dict(), 'model.pth') ``` 此代码将模型的参数保存到名为`model.pth`的文件中。 7. 加载模型: ``` model = MyModel() model.load_state_dict(torch.load('model.pth')) model.eval() ``` 通过加载参数文件,我们可以将模型的状态恢复到训练完成后的状态,然后使用`model.eval()`将模型设置为评估模式。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值