pytorch配置双显卡方式,使用双显卡跑代码!
pytorch配置双显卡方式,使用双显卡跑代码!
项目场景
Linux系统,pytorch环境
问题描述
使用的服务器有两张显卡,感觉一张显卡跑代码比较慢,想配置两张显卡同时跑代码,只需要在你的代码中添加几行,就可以使用双显卡,亲测有效。
解决方案
提示:这里填写该问题的具体解决方案:
先看以下官方示例代码,插入添加的地方是需要我们添加在代码中的代码行
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | importtorchimporttorch.nn as nnfromtorch.utils.data importDataset, DataLoaderimportos #######添加os.environ['CUDA_VISIBLE_DEVICES'] ='0,1'# 这里输入你的GPU_id# Parameters and DataLoadersinput_size =5output_size =2batch_size =30data_size =100#######添加device =torch.device("cuda:0"iftorch.cuda.is_available() else"cpu")# Dummy DataSetclassRandomDataset(Dataset):def__init__(self, size, length):self.len=lengthself.data =torch.randn(length, size)def__getitem__(self, index):returnself.data[index]def__len__(self):returnself.lenrand_loader =DataLoader(dataset=RandomDataset(input_size, data_size),batch_size=batch_size, shuffle=True)# Simple ModelclassModel(nn.Module):# Our modeldef__init__(self, input_size, output_size):super(Model, self).__init__()self.fc =nn.Linear(input_size, output_size)defforward(self, input):output =self.fc(input)print("\tIn Model: input size", input.size(),"output size", output.size())returnoutput################添加# Create Model and DataParallelmodel =Model(input_size, output_size)iftorch.cuda.device_count() > 1:print("Let's use", torch.cuda.device_count(), "GPUs!")model =nn.DataParallel(model)model.to(device)#Run the Modelfordata inrand_loader:input=data.to(device)output =model(input)print("Outside: input size", input.size(),"output_size", output.size()) | 
其中我将model = nn.DataParallel(model)修改为model = nn.DataParallel(model.cuda()),这一步直接参照网上修改的,因此这一步没有报错。
比如我自己在我代码中添加如下
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | frommodel.hash_model importDCMHT as DCMHTimportosfromtqdm importtqdmimporttorchimporttorch.nn as nnfromtorch.utils.data importDataLoaderimportscipy.io as sciofrom.base importTrainBasefrommodel.optimization importBertAdamfromutils importget_args, calc_neighbor, cosine_similarity, euclidean_similarityfromutils.calc_utils importcalc_map_k_matrix as calc_map_kfromdataset.dataloader importdataloader###############添加os.environ['CUDA_VISIBLE_DEVICES'] ='0,1'# 这里输入你的GPU_iddevice =torch.device("cuda:0"iftorch.cuda.is_available() else"cpu")classTrainer(TrainBase):def__init__(self,rank=0):args =get_args()super(Trainer, self).__init__(args, rank)self.logger.info("dataset len: {}".format(len(self.train_loader.dataset)))self.run()def_init_model(self):self.logger.info("init model.")linear =Falseifself.args.hash_layer =="linear":linear =Trueself.logger.info("ViT+GPT!")HashModel =DCMHTself.model =HashModel(outputDim=self.args.output_dim, clipPath=self.args.clip_path,writer=self.writer, logger=self.logger, is_train=self.args.is_train, linear=linear).to(self.rank)####################################添加self.model=nn.DataParallel(self.model.cuda())iftorch.cuda.device_count() >1:print("Lets use",torch.cuda.device_count(),"GPUs!")self.model.to(device)ifself.args.pretrained !="" andos.path.exists(self.args.pretrained):self.logger.info("load pretrained model.")self.model.load_state_dict(torch.load(self.args.pretrained, map_location=f"cuda:{self.rank}"))self.model.float()self.optimizer =BertAdam([{'params': self.model.clip.parameters(), 'lr': self.args.clip_lr},{'params': self.model.image_hash.parameters(), 'lr': self.args.lr},{'params': self.model.text_hash.parameters(), 'lr': self.args.lr}], lr=self.args.lr, warmup=self.args.warmup_proportion, schedule='warmup_cosine', b1=0.9, b2=0.98, e=1e-6, t_total=len(self.train_loader) *self.args.epochs,weight_decay=self.args.weight_decay, max_grad_norm=1.0)print(self.model) | 
添加以上代码后一般还会报如下错误
“AttributeError: ‘DataParallel’ object has no attribute ‘xxx’”
解决办法为先在dataparallel后的model调用module模块,然后再调用xxx
比如在上述我自己的代码中会报错
AttributeError: ‘DataParallel’ object has no attribute ‘clip’
解决办法:
是将model,修改为model.module.,后续报错大致相同,将你的代码中涉及到model.的地方修改为model.module.即可。
| 1 2 3 4 5 | self.optimizer =BertAdam([{'params': self.model.module.clip.parameters(), 'lr': self.args.clip_lr},{'params': self.model.module.image_hash.parameters(), 'lr': self.args.lr},{'params': self.model.module.text_hash.parameters(), 'lr': self.args.lr}], lr=self.args.lr, warmup=self.args.warmup_proportion,  | 
检查显卡使用情况
打开终端,在终端输入nvidia-smi命令可查看显卡使用情况

成功使用双显卡跑代码!
学习资料见知识星球。
以上就是今天要分享的技巧,你学会了吗?若有什么问题,欢迎在下方留言。
快来试试吧,小琥 my21ke007。获取 1000个免费 Excel模板福利!
更多技巧, www.excelbook.cn
欢迎 加入 零售创新 知识星球,知识星球主要以数据分析、报告分享、数据工具讨论为主;
1、价值上万元的专业的PPT报告模板。
2、专业案例分析和解读笔记。
3、实用的Excel、Word、PPT技巧。
4、VIP讨论群,共享资源。
5、优惠的会员商品。
6、一次付费只需129元,即可下载本站文章涉及的文件和软件。
