Pytorch 自动选择CPU或者GPU运行
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)
Pytorch单机多卡:
https://blog.csdn.net/xys430381_1/article/details/106635977
- PyTorch默认使用从0开始的GPU,且默认只使用0号GPU。如果要使用其他编号的GPU或者使用多块GPU,则要设置。
- pytorch并行后,假设batchsize设置为64,表示每张并行使用的GPU都使用batchsize=64来计算(单张卡使用时,使用batchsize=64比较合适时,多张卡并行时,batchsize仍为64比较合适,而不是64*并行卡数)。
- DataParallel 会自动拆分数据,并将作业订单发送到多个GPU上的多个模型。 在每个模型完成它们的工作之后,DataParallel 在将结果返回给你之前收集和合并结果。
有两种方法:
方法一:环境变量 + device + to(device)
第一步,指定备选的GPU
有如下两种方法来指定需要备选的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'
第二步,创建设备(device)
作用:将备选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")
第三步,将data和model放置到device上
# 如果他们已经在目标设备上则不会执行复制操作
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 即可
方法二:不建议用
Pytorch多机多卡:
参考
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