【Linux，Windows】指定GPU运行代码

查看GPU状态

nvidia-smi

程序中指定GPU

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'  # 这行代码须放在程序访问GPU之前，如程序的开头。

命令行指定GPU

CUDA_VISIBLE_DEVICES=0,1,2,3 python test.py  

多个GPU协同运行 

import os 
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

【切换】

export TRANSFORMERS_OFFLINE=1
export HF_DATASETS_OFFLINE=1
export CUDA_VISIBLE_DEVICES=1,2
export CUDA_DEVICE_ORDER=PCI_BUS_ID
nohup python data_train.py > log/log.txt 2>&1 &

同时在 cfg 文件中配置好参数

tt.arg.num_gpus = 1 if tt.arg_num_gpus is None else tt.arg.num_gpus  # 官方标准设置
# tt.arg.num_gpus = 2  # 或者根据自身机器舍子gpu的个数

在模型初始化函数中

# self.enc_module = enc_module.to(tt.arg.device)
# self.gnn_module = gnn_module.to(tt.arg.device)
self.enc_module = enc_module.cuda()
self.gnn_module = gnn_module.cuda()
​
if tt.arg.num_gpus > 1:
        print('Construct parallel model ...')
        self.enc_module = nn.DataParallel(self.enc_module)
        self.gnn_module = nn.DataParallel(self.gnn_module)
​
        print('done!\n')

pytorch中GPU加速测试

import torch
import time
from torch import autograd
#GPU加速
print(torch.__version__)
print(torch.cuda.is_available())

a=torch.randn(10000,1000)
b=torch.randn(1000,10000)
print(a)
print(b)
t0=time.time()
c=torch.matmul(a,b)
t1=time.time()

print(a.device,t1-t0,c.norm(2))

device=torch.device('cuda')
print(device)
a=a.to(device)
b=b.to(device)

t0=time.time()
c=torch.matmul(a,b)
t2=time.time()
print(a.device,t2-t0,c.norm(2))


t0=time.time()
c=torch.matmul(a,b)
t2=time.time()

print(a.device,t2-t0,c.norm(2))