insightface 切换GPU训练报错解决

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