解决多卡加载预训练模型时，卡0总会比其他卡多占用显存，多卡占用显存不均_多卡训练为什么占的内存多

作者：知新_RL | 2024-05-22 22:55:12

踩

多卡训练为什么占的内存多

出现问题：

再跑stdc官方给出代码时，使用多卡加载pretrained模型时，卡0上比卡123多三个进程，卡123都只有一个进程，中断训练后卡0全部进程都中断。即卡0在正常训练时比卡123多占用显存。（使用nvidia-smi命令查看）

这里发现问题存在于load模型的时候，直接load的话会导致参数加载到之前保存模型的device上（大部分情况下应该是只用cuda0去保存），这里可以将load函数加一个参数为map_location解决：

model_weights = torch.load(path, map_location='cpu')

改完之后发现仍旧存在这个问题

后来查看代码的时候发现，在网络的class内部，init_weight的时候也有一个load函数。这个函数内部无map_location参数。

因此，加入这个参数后如果卡0仍旧多占用显存和进程的时候，可能是某个隐藏的类别内部存在torch.load函数。多找找加上map_location参数即可使得全部的卡占用相同显存。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/知新_RL/article/detail/610273