1. Ver el dispositivo donde se encuentra el tensor:
data = data.cuda()#将数据转移到gpu上
print(data.device) # 输出:cuda:0
data = data.cpu()#将数据转移到cpu上
print(data.device) # 输出:cpu
2. Ver el dispositivo donde se encuentra el modelo
model = model.cuda()#将模型转移到gpu上
print(next(model.parameters()).device) # 输出:cuda:0
model = model.cpu()#将模型转移到cpu上
print(next(model.parameters()).device) # 输出:cpu
3. Hay dos formas comunes de cargar modelos y tensores en GPU en Pytorch.
Método 1:
# 如果GPU可用,将模型和张量加载到GPU上
if torch.cuda.is_available():
model = model.cuda()
x = x.cuda()
y = y.cuda()
Método 2:
# 分配到的GPU或CPU
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 将模型加到GPU
model=model.to(device)
# 将张量加到GPU
x=x.to(device)
y=y.to(device)
4. Especifique el código GPU
# 代码1:
torch.cuda.set_device(1)
# 代码2:
device = torch.device("cuda:1")
# 代码3:(官方推荐使用),
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
(如果你想同时调用两块GPU的话)
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'
Enlace de referencia: seleccione la GPU especificada en PyTorch
Tenga en cuenta que el código GPU especificado debe colocarse al principio del segmento del programa, como se muestra en la siguiente figura:
5. Ver la cantidad de GPU
torch.cuda.device_count()