insightface 切换GPU训练报错解决

举报
风吹稻花香 发表于 2021/06/08 23:34:57 2021/06/08
5.1k+ 0 0
【摘要】 insightface代码: https://github.com/deepinsight/insightface 在pytorch训练时,先用gpu1训练,再用gpu2训练,加载gpu1的预训练时报错, 解决方法: 先用cpu加载模型,再转换到新的gpu id上, 代码如下: self.weight: torch.Tensor = torch.load(sel...

insightface代码:

https://github.com/deepinsight/insightface

在pytorch训练时,先用gpu1训练,再用gpu2训练,加载gpu1的预训练时报错,

解决方法:

先用cpu加载模型,再转换到新的gpu id上,

代码如下:


       self.weight: torch.Tensor = torch.load(self.weight_name,map_location='cpu')
       self.weight=self.weight.cuda(local_rank)
       self.weight_mom: torch.Tensor = torch.load(self.weight_mom_name,map_location='cpu')
       self.weight_mom=self.weight_mom.cuda(local_rank)
  
 

partial_fc.py 完整代码如下:


      import logging
      import os
      import torch
      import torch.distributed as dist
      from torch.nn import Module
      from torch.nn.functional import normalize, linear
      from torch.nn.parameter import Parameter
      class PartialFC(Module):
      """
       Author: {Xiang
  
 

文章来源: blog.csdn.net,作者:网奇,版权归原作者所有,如需转载,请联系作者。

原文链接:blog.csdn.net/jacke121/article/details/117703681

【版权声明】本文为华为云社区用户转载文章,如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

作者其他文章

评论(0

抱歉,系统识别当前为高风险访问,暂不支持该操作

    全部回复

    上滑加载中

    设置昵称

    在此一键设置昵称,即可参与社区互动!

    *长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

    *长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。