https://blog.csdn.net/ssjdoudou/article/details/103640129

在程序最开始的地方加上

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

.cuda()  改为 .to(device)

Eg:

model.to(device)

input.to(device)

label.to(device)

https://blog.csdn.net/xys430381_1/article/details/106635977

有两种方法：

有如下两种方法来指定需要备选的GPU。

这些GPU 将在程序中可见（但不一定会使用，真正投入使用需要通过device()函数来再次选择和指定）

使用环境变量CUDA_VISIBLE_DEVICES的方式，类似于tensorflow指定GPU的方式（http://www.cnblogs.com/darkknightzh/p/6591923.html）。

直接终端中设定：

CUDA_VISIBLE_DEVICES=1

python代码中设定：

1、使用单卡

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "2"

2、使用多块卡的方法。

例如，使用0号和1号GPU’

import os

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'

作用：将备选GPU进一步选择和指定，真正投入使用中。

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

# 注意如果 device = torch.device("cuda")，则环境变量CUDA_VISIBLE_DEVICES中指定的全部GPU都会被拿来使用。

# 也可以通过 "cuda:0" 、"cuda:1"等指定环境变量CUDA_VISIBLE_DEVICES中指定的多块GPU中的某一块。

注意对应关系。例如：

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3, 4, 5"  # 将2, 3, 4, 5号GPU作为备选GPU

# 这样cuda:0表示使用 2 号GPU

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

# 如果他们已经在目标设备上则不会执行复制操作

input = data.to(device)

model = MyModule(...).to(device)

注意：如果有多个GPU，则model还需要多做一个操作（模型并行化）

第三步骤的多GPU版本如下：

input = data.to(device)

model = MyModule(...)

#模型并行化

if torch.cuda.device_count() > 1:

            print("Let's use", torch.cuda.device_count(), "GPUs!")

            model = nn.DataParallel(model)

model = model.to(device)

CUDA_VISIBLE_DEVICES=0,1,2,3 python  bert_pytorch.py  即可

方法二：不建议用

参考 

https://zhuanlan.zhihu.com/p/74792767

https://zhuanlan.zhihu.com/p/68717029

参考文献：

https://blog.csdn.net/ssjdoudou/article/details/103640129

https://blog.csdn.net/xys430381_1/article/details/106635977

https://zhuanlan.zhihu.com/p/74792767

https://zhuanlan.zhihu.com/p/68717029

ttps://zhuanlan.zhihu.com/p/86441879