Pytorch模型训练指定显卡

Pytorch模型训练指定显卡

参考资料：在pytorch中指定显卡

1、单卡训练

import torch
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'      #方式一：设置3号卡，并认为空间中就这张卡，查看gpu数量为1 # '0,1,2,3,4,5,6,7'
torch.cuda.set_device(1)		#⭐方式二：设置1号卡，识别当前空间的所有卡
CUDA_VISIBLE_DEVICES = 0 python3 train.py # 方式三：命令行中指定，不确定可行性

# 查看gpu是否可用
print('CUDA is available: {}'.format(torch.cuda.is_available())) # 方式一
use_gpu = torch.cuda.is_available() # 方式二
assert use_gpu, 'Current implementation does not support CPU mode. Enable CUDA.'

print('CUDA current_device: {}'.format(torch.cuda.current_device())) # ⭐查看当前使用gpu的索引
print('CUDA device_count: {}'.format(torch.cuda.device_count()))	# ⭐查看gpu数量
print(torch.cuda.get_device_name(0))    # 查看第1张gpu的名字

❓据说：os.environ[‘CUDA_VISIBLE_DEVICES’] 要定义在 import torch之前，否则会失效。如果引入了其他的.py文件中有 import torch ，且导入位置在 os.environ[‘CUDA_VISIBLE_DEVICES’] 之前同理会失效。

但是实际实现时发现都会失败。QuQ /home/workspace/wenqianli/yolo_v1_pytorch-master/train_yolo.py 在这个代码中，无法用 os.environ[‘CUDA_VISIBLE_DEVICES’] 设置显卡。但是，我记得在yolov3 中我有成功设置过，memo

后续有机会再测试

2、多卡训练

# 方式一
model = model.cuda() # 放第三行也可以
device_ids = [0, 1] 	# id为 0 和 1 的两块显卡
model = torch.nn.DataParallel(model, device_ids=device_ids)

# 方式二 model = model.cuda()放最后
device_ids = [0, 1]
model = torch.nn.DataParallel(model, device_ids=device_ids).cuda()

3、模型和数据加载到 GPU

.cpu() .cuda() .to(device)
据说.cuda()很少使用了

# 方式一
if torch.cuda.is_available():
    model = model.cuda()
    imgs,targets = imgs.cuda(),targets.cuda()

# 方式二
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
imgs, targets = imgs.to(device), targets.to(device)