pytorch gpu变成常用命令

官方文档：https://pytorch.org/docs/stable/cuda.html

0、设备命名：cpu，cuda:0(0代表gpu编号)，......，cuda:n

1、测试gpu是否可用:torch.cuda.is_available()

2、返回gpu可用数量：torch.cuda.device_count()

3、默认选择某一个gpu，若不存在，使用cpu：device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

4、将网络送入gpu运行：net.to(device)

5、将数据送入gpu运算：inputs, labels = inputs.to(device), labels.to(device)

6、返回index的gpu设备名字：torch.cuda.get_device_name(0)

7、返回当前设备索引：torch.cuda.current_device()

8、多gpu并行训练：

import torch
import torch.nn as nn

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

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


#如果具有多个gpu，使用数据并行技术
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)
#同样数据要放入设备中
inputs, labels = inputs.to(device), labels.to(device)

如果选择部分gpu使用，有两种方法：

i:直接使用DataParallel中device_ids来指定，device_ids=[0,1,2,3]

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

在正向传递中，模块在每个设备上复制，每个副本处理一部分输入。在向后传递期间，来自每个副本的渐变被加到原始模块中。

module：需要并行处理的模型
device_ids：并行处理的设备，默认使用所有的cuda
output_device：输出的位置，默认输出到cuda:0

ii:在shell脚本里面设置CUDA_VISIBLE_DEVICES变量

export CUDA_VISIBLE_DEVICES=0,1,2,3

python train.py

iii:在Python脚本里面设置gpu可见设备

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"