pytorch配置雙顯卡方式,使用雙顯卡跑代碼

更新時間：2024年06月26日 09:12:53 作者：好好好好飯

這篇文章主要介紹了pytorch配置雙顯卡方式,使用雙顯卡跑代碼,具有很好的參考價值,希望對大家有所幫助,如有錯誤或未考慮完全的地方,望不吝賜教

項目場景

Linux系統(tǒng)，pytorch環(huán)境

問題描述

使用的服務器有兩張顯卡，感覺一張顯卡跑代碼比較慢，想配置兩張顯卡同時跑代碼，只需要在你的代碼中添加幾行，就可以使用雙顯卡，親測有效。

解決方案

提示：這里填寫該問題的具體解決方案：

先看以下官方示例代碼，插入添加的地方是需要我們添加在代碼中的代碼行

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())，這一步直接參照網(wǎng)上修改的，因此這一步?jīng)]有報錯。

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

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.，后續(xù)報錯大致相同，將你的代碼中涉及到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,