[英]How to utilize all GPUs when dealing with pytorch code?
我有2个GPU,使用pytorch代码时,仅使用一个GPU。 我尝试了CUDA_VISIBLE_DEVICES=0,1 python xxx.py
,但是发生了
'CUDA_VISIBLE_DEVICES:找不到命令'
问题。 我还尝试在对象py文件中添加以下行:
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
但仍然只使用一个GPU。
您需要分别将训练数据并行化到每个GPU。 数据并行使用torch.nn.DataParallel
实现。 pytorch文档中的一个示例:
import torch
import torch.nn as nn
class DataParallelModel(nn.Module):
def __init__(self):
super().__init__()
self.block1 = nn.Linear(10, 20)
# wrap block2 in DataParallel
self.block2 = nn.Linear(20, 20)
self.block2 = nn.DataParallel(self.block2)
self.block3 = nn.Linear(20, 20)
def forward(self, x):
x = self.block1(x)
x = self.block2(x)
x = self.block3(x)
return x
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.