【pytorch】模型训练和推理使用的设备不一致导致的问题

最新推荐文章于 2024-08-12 11:05:45 发布
最新推荐文章于 2024-08-12 11:05:45 发布
阅读量1.7k 收藏 1
点赞数 1
CC 4.0 BY-SA版权
分类专栏： 机器学习&深度学习
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.youkuaiyun.com/JianJuly/article/details/107507448
当使用PyTorch在不同GPU间迁移已训练模型时，可能会遇到CUDA内存溢出错误。本文详细解释了错误产生的原因，即模型加载时默认指向被占用的显卡，以及如何通过设置map_location参数来解决此问题。

摘要生成于 C知道 ，由 DeepSeek-R1 满血版支持， 前往体验 >

问题

前几天使用0号显卡训练了模型，今天想要推理的时候，0卡被占用，所以选择使用1卡进行推理。
但是推理的时候却报错：RuntimeError: CUDA error: out of memory

原因

所有先前保存的模块，不论其设备，首先加载到CPU上，然后移动到之前保存它们的设备上（也就是0号显卡）。因为0号显卡被占用，所以显示out of memory

解决方法

在load模型的时候，添加map_location参数

torch.load(path_file, map_location='cuda:1')

参考资料

  1. https://github.com/apachecn/pytorch-doc-zh/blob/master/docs/1.0/jit.md
  2. pytorch在gpu训练加载模型,训练与测试在不同机器上遇到的问题总结
  3. pytorch将cpu训练好的模型参数load到gpu上，或者gpu->cpu上
如果我想要本地部署一个模型大小为 52GB 的推理服务，请问我需要什么样的GPU配置？请给我列一下详细的清单，报价方案,表格形式。
AI天才研究院
07-02 3959
要本地部署一个大小为 52GB 的推理服务，您需要一台具有足够GPU计算能力的服务器或工作站。请注意，这只是一个估算，价格可能会根据您所在地区和供应情况而有所不同。此外，还需要考虑到服务器的电源、散热、机箱等其他零部件和配件的成本。另外，部署和维护一个这样的系统需要专业知识和技能，建议寻求专业人士的帮助。
大模型的实践应用30-大模型训练和推理中分布式核心技术的应用
微学AI的博客
09-20 1144
在人工智能领域，近年来“大模型”（large language models, LLMs）已成为研究与应用的焦点，标志着AI技术进入了一个全新的发展阶段。以GPT系列、BERT、Turing-NLG等为代表的大型语言模型，在自然语言处理、文本生成、机器翻译等任务上展现了前所未有的能力，推动了所谓的“百模大战”。这场竞赛不仅限于学术界，科技巨头如Google、Microsoft、阿里云等也纷纷加入，不断刷新模型参数量的纪录，力求在人工智能的军备竞赛中占据领先地位。
模型推理时CPU和GPU结果不一致
小松的博客
09-27 1155
这个设置是默认为True的，会让程序在开始时花费一点额外时间，为整个网络的每个卷积层搜索最适合它的卷积实现算法，进而实现网络的加速。也就是说训练时如果使用的最优的卷积算法，而关闭掉后选择了别的卷积算法，就会导致结果不一致。
pytorch连续多次推理同样的数据，结果居然不一样？
cs_softwore的专栏
03-09 4141
深度学习中经常会用到反卷积操作，在使用pytorch时，发现如下问题，特做记录。 语句1：self.deconv1 = nn.ConvTranspose2d(in_channel, out_channel, kernel_size=4, stride=2, padding=1， output_padding=0， bias=False) 语句2：self.deconv1 = nn.Conv...
torch load或者load_state_dict（）后模型不相同
Yonggie的博客
05-01 8479
一言以蔽之 把原先模型eval，把加载后模型eval。才保证两模型完全一样。 保存模型 torch.save(net,'./model.pth') 或者 torch.save(net.state_dict(),'./model-dict.pth') 加载模型 net=torch.load('./model.pth') 或者 net.load_state_dict(torch.load('./mo...
讨论PyTorch中模型加载时参数不一致的情况
敲代码的小风
03-13 8030
参考链接: torch.nn.Module.load_state_dict(state_dict, strict=True) 首先保存参数不同的模型: 保存模型的代码: # 分别为来自CPU和GPU的模型保存参数， # 并再次将模型的参数加载到位于CPU或者GPU的模型上 import torch import torch.nn as nn torch.manual_seed(seed=20200910) class Model(torch.nn.Module): def __init__(se
大模型训练和推理
weixin_45325331的博客
02-27 4665
大模型训练moe是指混合专家（Mixture-of-experts，简称MoE）的，这是一种用于提高大型神经网络性能的技术，它可以将模型的一部分替换为多个专家，每个专家只负责处理一部分输入数据，从而增加模型的参数量和表达能力，同时减少计算量和训练时间。指令可以作为模型的输入的一部分，也可以作为模型的输出的一部分，具体取决于任务的类型。上下文扩展技术的，这是一种用于提高大型语言模型（LLM）对长文本的处理能力的技术，它可以让模型在不增加参数量的情况下，扩展其注意力范围，从而提高模型的语言理解和生成能力。
踩坑日记：pytorch模型训练train模式和eval模式结果相差甚远？
weixin_39353767的博客
02-07 3910
如我在文中对相关术语和概念的描述错误或不准确，欢迎评论区指正！谢谢～ 最近在利用pytorch进行transformer模型训练的时候发现，明明训练的时候loss已经降到非常低、在调试模式下结果也十分理想，在测试集上生成的结果却非常离谱， ...
【图像去噪】论文复现：新手入门必看！DnCNN的Pytorch源码训练测试全流程解析！为源码做详细注释！补充DnCNN-B和DnCNN-3的模型训练和测试！附各种情况下训练好的模型权重文件！
热门推荐
主要更新底层视觉（去噪、超分等）相关的科研内容，形式为【论文精读】+【论文复现】
08-12 1万+
- 以官方Pytorch源代码为基础，在DnCNN-S的基础上，增添DnCNN-B/CDnCNN-B、DnCNN-3模型训练和测试复现，代码注释非常详细，无论是科研还是应用，新手小白都能看懂，学习阅读毫无压力，去噪入门必备，适用于去噪、超分、JPEG去块任务； - 提供新增后的完整代码和训练好的模型权重文件，模型性能与论文中近似，可不训练直接测试； - 理论和源码结合，进一步加深理解算法原理、明确训练和测试流程 - 更换路径和相关参数即可训练自己的图像数据集
pytorch调用不了多个gpu_使用Pytorch在多GPU下保存和加载训练模型参数遇到的问题...
weixin_26824299的博客
12-24 441
最近使用Pytorch在学习一个深度学习项目，在模型保存和加载过程中遇到了问题，最终通过在网卡查找资料得已解决，故以此记之，以备忘却。首先，是在使用多GPU进行模型训练的过程中，在保存模型参数时，应该使用类似如下代码进行保存：torch.save({'epoch':epoch,'state_dict':model.module.state_dict(),'optimizer':optimiz...
深度学习在训练和测试阶段使用显卡的情况是否必须完全一致？
无敌三脚猫的专栏
04-24 1614
当然，在使用GPU进行推理时，由于GPU的并行计算能力较强，可以显著提高推理速度，特别是在处理大型数据集或者复杂模型时更加明显。因此，如果您有可用的GPU资源，建议在测试过程中使用GPU来进行推理，以获得更快的速度和更好的性能。深度学习模型在进行训练时采用多张显卡进行训练，测试时是不是就与显卡无关了，也就是说可以利用CPU做推理，也可以使用GPU做推理。是的，在深度学习模型训练时采用多张显卡进行训练，测试时模型的预测过程与显卡无关。而在测试阶段，模型的预测过程则不需要如此大量的计算资源，因此可以使用。
tensor不在同一设备上：创建tensor后就马上.to(device)
巨无霸汉堡加大的博客
04-03 1097
创建tensor就要.to(device) labels = torch.zeros( logits.shape[0], logits.shape[1] ).to(device).scatter_( dim=1, index=labels, src=torch.ones( labels.shape[0], labels.shape[1] ).to(device) )
训练AI数据模型所需要的高性能计算机配置
polsnet的专栏
04-04 1万+
AI模型训练是一种机器学习的过程，通过训练深度学习模型来自动化处理数据和完成任务。AI训练可以帮助企业和研究人员开发出更加智能、高效的应用，从而提高生产力和创新能力。
yolov5的GPU更换，训练模型和使用
qq_37666407的博客
03-24 3064
最近正在研究的一些项目都需要使用物体识别这个功能，所以专门研究了一下yolov5的使用，讲的可能不是很精确，但是还是值得记录一下的。
mindspore，resnet双卡训练报错
weixin_45666880的博客
06-22 302
问题描述：飞腾+300T训练卡，麒麟V10 SP1操作系统mindspore1.6模型：resnet50问题现象：单卡训练没问题，双卡训练报错 解答：这种错误的原因是该卡在最大等待时长内，没有收到其他卡的同步消息。 而没有收到其他卡的同步消息，有很多种可能： 其他卡停了； 其他卡还在执行编译，且编译时间略长，超出了该卡的最大等待时长； 对于第一种情况，需要检查是否存在某张卡有问题无法正常使用的情况，可能是硬件问题，也可能是软件问题，可以通过查看每张卡的plog得知进一步的信息。...
【MindSpore】【分布式训练】Atlas800-9000裸机单机多卡训练modelzoo模型报错
weixin_45666880的博客
08-21 658
export ASCEND_SLOG_PRINT_TO_STDOUT=1，配置这个环境变量才能看到真实的底层报错。3、使用models/hccl_tools.py生成对应的rank_table.json文件。4、使用0,1卡进行单机多卡训练（2卡训练和8卡训练报错一致，方便起见使用2卡），报错。1、配置完环境后从MindSpore官方modelzoo拉取resnet代码。2、单机单卡训练正常（run_standalone_train.sh）。5、尝试多台atlas800设备，均存在相同错误。
keras相同模型，相同输入，不同机器输出完全不同的原因！！
Avalon
11-12 5395
问题关键：keras版本导致，不同机器keras版本不一致，将不同机器的keras版本保持一致就行了...  最近在做一个以图搜图的项目，想到可以用keras的预训练模型提取特征，然后再计算各个图片特征的余弦距离。自己有差不多10万张图片，一开始在GPU服务器上提特征，然后将特征写在数据库里，在服务器上模拟识别效果还不错。然后自己在本地的电脑上模拟识别，先利用预训练模型提取特征然后再与GPU提取...
关于pytorch tensor/params没有放在同一个设备上的经验
yz1019117968的博客
01-23 3343
前言 用pytorch做深度学习时候，全部tensor和变量都要放到同一个设备上, i.e. GPU or CPU。否则会报一下错误： RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_mm) 出现以上bug，目前我遇到的
如何在 PyTorch 中加载预训练模型权重并直接用于推理
最新发布
04-28
<think>好的，我现在需要解决用户的问题，即如何在PyTorch中加载预训练模型权重并用于推理。根据用户提供的引用内容，我需要综合这些信息来构建一个清晰的步骤指南。 首先，用户提到的引用[1]提到使用torch.load(PATH)加载模型，并调用model.eval()进入推理模式。这应该是基础步骤。不过，用户可能需要注意模型结构的匹配问题，如果预训练模型的架构与当前模型不完全一致，可能需要调整加载方式。 引用[2]和[3]讨论了迁移学习中的微调，特别是如何加载部分参数并冻结某些层。虽然用户主要关注推理，但了解如何正确加载权重是关键。例如，引用[3]中的代码展示了如何处理模型参数不匹配的情况，比如忽略分类层或处理BN层的额外参数。这可能对用户有帮助，尤其是在模型结构有变化时。 引用[4]提到的问题是关于加载相同结构部分但输出不一致的情况，可能是因为某些层未被正确固定。虽然用户的问题是关于推理，但正确加载权重并确保模型状态正确（如eval模式）可能影响输出结果，这点需要注意。 接下来，我需要将这些信息整合成一个步骤化的教程。首先加载模型结构，然后加载权重，处理可能的键不匹配问题，设置模型为评估模式，并进行前向传播。同时，需要提到常见的错误，如忘记调用eval()导致dropout或BN层行为不一致，或者结构不匹配时的解决方法。 可能还需要考虑使用torch.no_grad()来减少内存消耗，并加快计算速度。此外，示例代码将帮助用户更好地理解流程，比如加载预训练的ResNet并进行推理。 最后，根据用户的要求，生成相关问题，这些问题应围绕加载模型的不同情况、处理结构差异、性能优化等。需要确保问题覆盖用户可能遇到的后续疑问。</think>### PyTorch加载预训练模型权重进行推理的完整教程 #### 一、基础加载流程（结构完全匹配） 1. **定义模型结构** 确保模型类定义与预训练权重完全一致，包括层名称、参数顺序、网络分支结构。 2. **加载权重与模式切换** ```python import torch from models import YourModelClass # 替换为实际模型定义 # 初始化模型 model = YourModelClass() # 加载预训练权重（推荐使用map_location参数确保设备兼容） checkpoint = torch.load('pretrained.pth', map_location='cpu') model.load_state_dict(checkpoint['model_state_dict']) # 根据实际键名调整 # 切换推理模式（关键步骤） model.eval() ``` #### 二、常见场景解决方案 ▌**场景1：部分结构不匹配**（参考引用[3][^3]） ```python # 步骤1：打印当前模型与预训练模型的键名差异 pretrained_dict = torch.load('vgg16_bn.pth') model_dict = model.state_dict() # 步骤2：筛选可加载参数（忽略无关层） filtered_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape} # 步骤3：特殊处理BN层参数（处理num_batches_tracked字段） if 'num_batches_tracked' in model_dict.keys(): filtered_dict = {k: v for k, v in filtered_dict.items() if 'num_batches_tracked' not in k} # 步骤4：加载并设置严格模式 model.load_state_dict(filtered_dict, strict=False) ``` ▌**场景2：仅需加载部分组件**（参考引用[4][^4]） ```python # 针对backbone模块加载 backbone_dict = {k.replace('backbone.', ''): v for k, v in checkpoint.items() if k.startswith('backbone')} model.backbone.load_state_dict(backbone_dict) # 冻结指定层参数 for param in model.backbone.parameters(): param.requires_grad_(False) ``` #### 三、执行推理的标准流程 ```python # 数据预处理规范 from torchvision import transforms preprocess = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 推理过程封装 with torch.no_grad(): # 禁用梯度计算 input_tensor = preprocess(raw_image).unsqueeze(0) if torch.cuda.is_available(): input_tensor = input_tensor.to('cuda') model.to('cuda') outputs = model(input_tensor) probabilities = torch.nn.functional.softmax(outputs[0], dim=0) ``` #### 四、关键注意事项 1. **模式区别** `model.eval()`会关闭dropout层和BN层的训练模式统计量计算，这对输出一致性至关重要[^1] 2. **设备管理** 通过`map_location`参数实现跨设备加载： ```python # GPU加载CPU训练的模型 torch.load('cpu_model.pth', map_location=torch.device('cuda:0')) ``` 3. **版本兼容** 当出现`Unexpected key(s) in state_dict`警告时，可使用`strict=False`参数跳过不匹配层（需验证关键层是否加载成功） #### 五、完整示例（以ResNet50为例） ```python import torchvision.models as models # 初始化模型 model = models.resnet50(pretrained=False) # 加载自定义权重 checkpoint = torch.load('custom_resnet50.pth') model.load_state_dict(checkpoint) # 准备输入 model.eval() input_tensor = torch.randn(1, 3, 224, 224) # 执行推理 with torch.no_grad(): output = model(input_tensor) print(f"Output shape: {output.shape}") ```
评论 4
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

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

抵扣说明：

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

余额充值