model
先把model都定义好,再统一放到GPU上,不然可能会报错
不要动optimizer
if len(conf.device_ids)>1:# 多GPU
self.model = torch.nn.DataParallel(self.model, device_ids=conf.device_ids).cuda()
self.head = torch.nn.DataParallel(self.head, device_ids=conf.device_ids).cuda()
#self.optimizer = torch.nn.DataParallel(self.optimizer, device_ids=conf.device_ids).cuda()
else: # 单GPU
self.model = self.model.to(conf.device)
self.head = self.head.to(conf.device)
#self.optimizer = self.optimizer.to(conf.device)
data
imgs = imgs.cuda()
labels = labels.cuda()
一些BUG
有些帖子说在取用state_dict
和param_groups
之前要加入.module,我没有改动这一块。