Pytorch指定GPU的操作小结

type
status
date
slug
summary
tags
category
icon
password
comment_flag
SLUGS
转移tensor操作等价表：
Title
.to(name)
.to(device)
.cuda()
GPU
to('cpu')
to(torch.device('cpu'))
cpu()
Current GPU
to('cuda')
to(torch.device('cuda'))
cuda()
Specific GPU
to('cuda:1')
to(torch.device('cuda:1'))
cuda(device=1)
第二行的操作是将tensor放到current cuda device上，当前cuda设备默认为"cuda:0"（即gpu_list[0]），可以通过torch.cuda.set_device()重新指定current cuda device：
 
对于pytorch的深度学习代码，main.py里通常是：
说明：
  1. os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"似乎和import torch顺序可以换，但建议前者至少得在torch.device("cuda:0" if torch.cuda.is_available() else "cpu")之前。
  1. os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list定义了程序可见gpu的列表（不定义该环境变量则默认全部可见），即能被device_x=torch.device("cuda:x" if torch.cuda.is_available() else "cpu")所指定的gpu的列表。后者语句中的device_x其实是卡号为gpu_list[x]的设备。
  1. os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list中的gpu_list建议为实际要用的卡号列表。如果gpu_list中含有程序实际不用的卡号，当前程序也会在实际不用的卡号后面占0M的显存。所以，os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list是建议一定要定义的。
 
以上内容若有纰漏，欢迎斧正。

参考

  1. https://stackoverflow.com/questions/62907815/pytorch-what-is-the-difference-between-tensor-cuda-and-tensor-totorch-device
  1. https://discuss.pytorch.org/t/what-is-the-difference-between-using-tensor-cuda-and-tensor-to-torch-device-cuda-0/89259
Loading...