Pytorch模型加载详解

小蓝 • 2025-01-03 14:49 • 编程

一、Pytorch模型加载概述

PyTorch是一个使用GPU和CPU优化的深度学习张量库，它也是一个动态神经网络构建工具。Pytorch模型加载是将已训练好的模型加载到内存中，以便使用。模型加载是模型应用的前提。Pytorch模型加载涉及到模型的序列化，反序列化和模型参数的赋值等操作。Pytorch中支持多种不同的序列化和反序列化方法，包括pickle和h5py等，其中最常用的方法是torch.save()和torch.load()函数。

二、Pytorch模型的保存与加载

1、保存模型


import torch

# 定义模型
model = torch.nn.Sequential(
   torch.nn.Linear(10, 100),
   torch.nn.ReLU(),
   torch.nn.Linear(100, 1),
   torch.nn.Sigmoid()
)

# 保存模型
torch.save(model.state_dict(), 'model.pth')

在上面的代码中，我们首先定义了一个简单的模型，然后使用torch.save()函数将模型参数保存到了’model.pth’文件中。

2、加载模型


import torch

# 定义模型
model = torch.nn.Sequential(
   torch.nn.Linear(10, 100),
   torch.nn.ReLU(),
   torch.nn.Linear(100, 1),
   torch.nn.Sigmoid()
)

# 加载模型
model.load_state_dict(torch.load('model.pth'))

在上面的代码中，我们首先定义了一个模型，然后使用torch.load()函数加载’model.pth’文件中的参数，最后使用load_state_dict()函数将参数赋值给模型。

三、Pytorch模型加载的不同形式

1、加载整个模型


import torch

# 保存模型
torch.save(model, 'model.pth')

# 加载模型
model = torch.load('model.pth')

在这个例子中，我们使用了torch.save()函数保存整个模型，并使用torch.load()加载整个模型。

2、多个模型的保存与加载


import torch

# 定义多个模型
model1 = torch.nn.Sequential(
   torch.nn.Linear(10, 100),
   torch.nn.ReLU(),
   torch.nn.Linear(100, 1),
   torch.nn.Sigmoid()
)

model2 = torch.nn.Sequential(
   torch.nn.Linear(10, 100),
   torch.nn.ReLU(),
   torch.nn.Linear(100, 10),
   torch.nn.Softmax()
)

# 保存多个模型
torch.save({
   'model1': model1.state_dict(),
   'model2': model2.state_dict()
}, 'multi_model.pth')

# 加载多个模型
checkpoint = torch.load('multi_model.pth')
model1.load_state_dict(checkpoint['model1'])
model2.load_state_dict(checkpoint['model2'])

在这个例子中，我们定义了两个模型，然后使用torch.save()函数保存多个模型的参数，并使用torch.load()加载多个模型的参数，最后使用load_state_dict()函数将参数赋值给对应的模型。

3、CPU/GPU间的模型加载


import torch

# 保存模型
torch.save(model.state_dict(), 'model.pth')

# 在CPU中加载模型
model_cpu = torch.nn.Sequential(
   torch.nn.Linear(10, 100),
   torch.nn.ReLU(),
   torch.nn.Linear(100, 1),
   torch.nn.Sigmoid()
)
model_cpu.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))

# 在GPU中加载模型
model_gpu = torch.nn.Sequential(
   torch.nn.Linear(10, 100),
   torch.nn.ReLU(),
   torch.nn.Linear(100, 1),
   torch.nn.Sigmoid()
).to('cuda')
model_gpu.load_state_dict(torch.load('model.pth'))

在这个例子中，我们首先使用torch.save()函数保存模型参数，然后使用torch.load()函数在CPU和GPU中加载模型。需要注意的是，在加载模型时需要使用map_location参数将模型参数映射到对应的设备上。如果我们想要将模型加载到GPU上，则需要通过.to(‘cuda’)将模型转移到GPU上。

四、总结

PyTorch模型加载是将已训练好的模型加载到内存中，以便进行推理或微调。在PyTorch中，我们可以使用torch.save()函数保存模型权重，使用torch.load()函数加载模型权重，并使用load_state_dict()函数将权重赋值给模型。同时，我们还可以保存和加载多个模型，将模型加载到不同的设备上运行。

原创文章，作者：小蓝，如若转载，请注明出处：https://www.506064.com/n/308407.html