目录
  • 项目场景
  • 问题描述
  • 解决方案
    • 比如我自己在我代码中添加如下
    • 添加以上代码后一般还会报如下错误
    • 比如在上述我自己的代码中会报错
    • 检查显卡使用情况
  • 总结

    项目场景

    Linux系统,pytorch环境

    问题描述

    使用的服务器有两张显卡,感觉一张显卡跑代码比较慢,想配置两张显卡同时跑代码,只需要在你的代码中添加几,就可以使用双显卡,亲测有效。

    解决方案

    提示:这里填写该问题的具体解决方案:

    先看以下官方示例代码,插入添加的地方是需要我们添加在代码中的代码行

    import torch
    import torch.nn as nn
    from torch.utils.data import Dataset, DataLoader
    import os 
    #######添加
    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' # 这里输入你的GPU_id
     
    # Parameters and DataLoaders
    input_size = 5
    output_size = 2
     
    batch_size = 30
    data_size = 100
    #######添加
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     
    # Dummy DataSet
    class RandomDataset(Dataset):
     
        def __init__(self, size, length):
            self.len = length
            self.data = torch.randn(length, size)
     
        def __getitem__(self, index):
            return self.data[index]
     
        def __len__(self):
            return self.len
     
    rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
                             batch_size=batch_size, shuffle=True)
     
    # Simple Model
    class Model(nn.Module):
        # Our model
     
        def __init__(self, input_size, output_size):
            super(Model, self).__init__()
            self.fc = nn.Linear(input_size, output_size)
     
        def forward(self, input):
            output = self.fc(input)
            print("\tIn Model: input size", input.size(),
                  "output size", output.size())
     
            return output
    ################添加
    # Create Model and DataParallel
    model = Model(input_size, output_size)
    if torch.cuda.device_count() > 1:
      print("Let's use", torch.cuda.device_count(), "GPUs!")
      model = nn.DataParallel(model)
    model.to(device)
     
     
    #Run the Model
    for data in rand_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()),这一步直接参照网上修改的,因此这一步没有报错。

    比如我自己在我代码中添加如下

    from model.hash_model import DCMHT as DCMHT
    import os
    from tqdm import tqdm
    import torch
    import torch.nn as nn
    from torch.utils.data import DataLoader
    import scipy.io as scio
     
     
    from .base import TrainBase
    from model.optimization import BertAdam
    from utils import get_args, calc_neighbor, cosine_similarity, euclidean_similarity
    from utils.calc_utils import calc_map_k_matrix as calc_map_k
    from dataset.dataloader import dataloader
    ###############添加
    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' # 这里输入你的GPU_id
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     
    class Trainer(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 = False
            if self.args.hash_layer == "linear":
                linear = True
     
            self.logger.info("ViT+GPT!")
            HashModel = DCMHT
            self.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())
            if torch.cuda.device_count() >1:
                print("Lets use",torch.cuda.device_count(),"GPUs!")
            self.model.to(device)
     
            if self.args.pretrained != "" and os.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.即可。

    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命令可查看显卡使用情况

    成功使用双显卡跑代码!

    总结

    以上为个人经验,希望能给大家一个参考,也希望大家多多支持小闻网。

    声明:本站所有文章,如无特殊说明或标注,均为本站原创发布。任何个人或组织,在未征得本站同意时,禁止复制、盗用、采集、发布本站内容到任何网站、书籍等各类媒体平台。如若本站内容侵犯了原著者的合法权益,可联系我们进行处理。