python中torch.load中的map_location参数使用

 更新时间:2024年03月18日 09:30:00   作者:高斯小哥  
在PyTorch中,torch.load()函数是用于加载保存模型或张量数据的重要工具,map_location参数为我们提供了极大的灵活性,具有一定的参考价值,感兴趣的可以了解一下

引言

在PyTorch中,torch.load()函数是用于加载保存模型或张量数据的重要工具。当我们训练好一个深度学习模型后,通常需要将模型的参数(或称为状态字典,state_dict)保存下来,以便后续进行模型评估、继续训练或部署到其他环境中。在加载这些保存的数据时,map_location参数为我们提供了极大的灵活性,以决定这些数据应该被加载到哪个设备上。本文将详细解析map_location参数的功能和使用方法,并通过实战案例来展示其在不同场景下的应用。

map_location参数详解

map_location参数在torch.load()函数中扮演着至关重要的角色。它决定了从保存的文件中加载数据时应将它们映射到哪个设备上。在PyTorch中,设备可以是CPU或GPU,而GPU可以有多个,每个都有其独立的索引。map_location的灵活使用能够让我们轻松地在不同设备之间迁移模型,从而充分利用不同设备的计算优势。

map_location参数的数据类型

map_location参数的数据类型可以是:

参数类型描述示例
字符串(str)预定义的设备字符串,指定目标设备。1. 'cpu':加载到CPU上;
2. 'cuda:X':加载到索引为X的GPU上。
torch.device对象一个表示目标设备的torch.device对象。1.torch.device('cpu'):加载到CPU上;
2. torch.device('cuda:1'):加载到索引为1的GPU上。
可调用对象(callable)一个接收存储路径并返回新位置的函数。lambda storage, loc: storage.cuda(1):将每个存储对象移动到索引为1的GPU上。
字典(dict)一个将存储路径映射到新位置的字典。{'cuda:1':'cuda:0'}:将原本在GPU 1上的张量加载到GPU 0上。

map_location参数的使用场景

  • CPU加载:当你想在CPU上加载模型时,可以设置map_location='cpu'。这适用于那些不需要GPU加速的推理任务,或者在没有GPU的环境中部署模型。

  • 指定GPU加载:如果你有多个GPU,并且想将模型加载到特定的GPU上,可以使用'cuda:X'格式的字符串,其中X是GPU的索引。这在多GPU环境中非常有用,可以确保模型加载到指定的设备上。

  • 自动选择GPU:如果你只想在GPU上加载模型,但不关心具体是哪一个GPU,可以设置map_location=torch.device('cuda')。这会自动选择第一个可用的GPU来加载模型。

  • 保持原始设备:如果你想保持模型在加载时的原始设备(即如果模型原先是在GPU上训练的,就仍然在GPU上加载;如果是在CPU上,就在CPU上加载),可以使用map_location=Nonemap_location=torch.device('cpu')(对于CPU模型)和map_location=torch.device('cuda')(对于GPU模型)。

  • 自定义映射逻辑:通过传递一个可调用对象,你可以实现更复杂的映射逻辑。例如,你可以编写一个函数,根据存储路径或模型结构来决定将模型加载到哪个设备上。这在需要根据特定条件动态选择加载设备时非常有用。

代码实战(详细注释)

下面将通过几个实战案例来展示map_location参数在不同场景下的应用。

案例1:从文件加载张量到CPU

# 案例1:从文件加载张量到CPU
# 使用torch.load()函数加载tensors.pt文件中的所有张量到CPU上
tensors = torch.load('tensors.pt')

案例2:指定设备加载张量

# 案例2:指定设备加载张量
# 使用torch.load()函数并指定map_location参数为CPU设备,加载tensors.pt文件中的所有张量到CPU上
tensors_on_cpu = torch.load('tensors.pt', map_location=torch.device('cpu'))

案例3:使用匿名函数指定加载位置

# 案例3:使用函数指定加载位置
# 使用torch.load()函数和map_location参数为一个lambda函数,该函数不做任何改变,保持张量原始位置(通常是CPU)
tensors_original_location = torch.load('tensors.pt', map_location=lambda storage, loc: storage)

案例4:将张量加载到指定GPU

# 案例4:将张量加载到指定GPU
# 使用torch.load()函数和map_location参数为一个lambda函数,该函数将张量移动到索引为1的GPU上
tensors_on_gpu1 = torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))

案例5:张量从一个GPU映射到另一个GPU

# 案例5:张量从一个GPU映射到另一个GPU
# 使用torch.load()函数和map_location参数为一个字典,将原本在GPU 1上的张量映射到GPU 0上
tensors_mapped = torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

案例6:从io.BytesIO对象加载张量

# 案例6:从io.BytesIO对象加载张量
# 打开tensor.pt文件并读取内容到BytesIO缓冲区
with open('tensor.pt', 'rb') as f:
    buffer = io.BytesIO(f.read())
    
# 使用torch.load()函数从BytesIO缓冲区加载张量
tensors_from_buffer = torch.load(buffer)

案例7:使用ASCII编码加载模块

# 案例7:使用ASCII编码加载模块
# 使用torch.load()函数和encoding参数为'ascii',加载module.pt文件中的模块(如神经网络模型)
model = torch.load('module.pt', encoding='ascii')

这些案例代码和注释展示了如何使用torch.load()函数的不同map_location参数和编码设置来加载张量和模型。这些设置对于控制数据加载的位置和格式非常重要,特别是在跨设备或跨平台加载数据时。

参考文档

[1] PyTorch官方文档

到此这篇关于python中torch.load中的map_location参数使用的文章就介绍到这了,更多相关python torch.load map_location参数内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家!

相关文章

  • selenium自动化测试简单准备

    selenium自动化测试简单准备

    本文主要介绍了selenium自动化测试简单准备,文中通过示例代码介绍的非常详细,具有一定的参考价值,感兴趣的小伙伴们可以参考一下
    2022-01-01
  • python实现csdn全部博文下载并转PDF

    python实现csdn全部博文下载并转PDF

    我们学习编程,在学习的时候,会有想把有用的知识点保存下来,我们可以把知识点的内容爬下来转变成pdf格式,方便我们拿手机可以闲时翻看,是很方便的,本文就介绍一下如何实现
    2021-06-06
  • Python如何实现SSH远程连接与文件传输

    Python如何实现SSH远程连接与文件传输

    这篇文章主要介绍了Python如何实现SSH远程连接与文件传输问题,具有很好的参考价值,希望对大家有所帮助。如有错误或未考虑完全的地方,望不吝赐教
    2023-05-05
  • Python全栈之文件函数和函数参数

    Python全栈之文件函数和函数参数

    这篇文章主要为大家介绍了Python的文件函数和函数参数,具有一定的参考价值,感兴趣的小伙伴们可以参考一下,希望能够给你带来帮助
    2021-12-12
  • Window版下在Jupyter中编写TensorFlow的环境搭建

    Window版下在Jupyter中编写TensorFlow的环境搭建

    这篇文章主要介绍了Window版下在Jupyter中编写TensorFlow的环境搭建,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧
    2020-04-04
  • Python采集王者最低战力信息实战示例

    Python采集王者最低战力信息实战示例

    这篇文章主要为大家介绍了Python采集王者最低战力信息实战示例解析,有需要的朋友可以借鉴参考下,希望能够有所帮助,祝大家多多进步,早日升职加薪
    2023-04-04
  • Python中使用Lambda函数的5种用法

    Python中使用Lambda函数的5种用法

    这篇文章主要介绍了Python中使用Lambda函数的5种用法,本文给大家介绍的非常详细,对大家的学习或工作具有一定的参考借鉴价值,需要的朋友可以参考下
    2021-04-04
  • Python中三种花式打印的示例详解

    Python中三种花式打印的示例详解

    在Python中有很多好玩的花式打印,我们今天就来挑战下面三个常见的花式打印。文中的示例代码讲解详细,感兴趣的小伙伴快跟随小编一起学习一下吧
    2022-03-03
  • 图文详解在Anaconda安装Pytorch的详细步骤

    图文详解在Anaconda安装Pytorch的详细步骤

    Anaconda指的是一个开源的Python发行版本,其包含了conda、Python等180多个科学包及其依赖项,下面这篇文章主要给大家介绍了关于在Anaconda安装Pytorch的详细步骤,需要的朋友可以参考下
    2022-07-07
  • python多进程并发demo实例解析

    python多进程并发demo实例解析

    这篇文章主要介绍了python多进程并发demo实例解析,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下
    2019-12-12

最新评论