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