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())

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)

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,