解决多卡加载预训练模型0卡内存占用大的问题

最新推荐文章于 2024-12-04 23:13:48 发布

_张一凡

最新推荐文章于 2024-12-04 23:13:48 发布

阅读量1.2k

点赞数 1

文章标签：深度学习人工智能

本文链接：https://blog.youkuaiyun.com/weixin_31588979/article/details/129126497

版权

在进行多卡训练模型时遇到显存不足的错误，特别是第一张卡占用显著。文章提到，通过将预训练模型的参数加载到CPU上而非直接在GPU上，可以缓解这一问题，从而避免显存溢出。修改后的代码示例展示了如何使用`torch.load`时指定`map_location=cpu`来实现这一方法。

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

一、问题描述

在多卡加载与训练模型的时候，出现显卡内存不足的错误，batchsize放到很小以后观察发现第一张卡占用内存比其他几张大了很多。

在这里插入图片描述

     #迁移学习
    if cfg.feature_extract:
        #print("feature_extract")
        backbone_weight_path=cfg.backbone_weight_path

        load_model = torch.load(backbone_weight_path)
        backbone.load_state_dict(load_model, strict=False)

解决方法：
把预训练模型参数map到cpu上去。

修改代码

     #迁移学习
    if cfg.feature_extract:
        #print("feature_extract")
        backbone_weight_path=cfg.backbone_weight_path

        load_model = torch.load(backbone_weight_path,map_location='cpu')
        backbone.load_state_dict(load_model, strict=False)