设备指定_指定GPU设备常用操作/cuda设备指定/查看模型在哪

x.0 查看有哪些设备

当你的model网络模型和张量数据tensor在CPU上时，其实是CPU和内存配合工作；而当你的model网络模型和张量数据tensor在GPU上时，其实是对应的GPU和对应的GPU的显存配合工作。

下面方法查看你有哪些设备，

print(torch.device("cpu"))
print(torch.device("cuda:1"))
print(torch.cuda.is_available())
print(torch.cuda.device_count())

x.1 指定哪些cuda设备可见

指定GPU设备常用操作，代码见下：

'''
设备指定

在跑多GPU时候，要确定你要计算的数据在同一张卡上。
'''

"""
(1) 指定设备

这个地方要在os.environ环境变量中指定CUDA可见参数的原因在于: 如果你不指定，则Pytorch会在其余可见的显卡上占用2~3MB的显存
"""
# 指定设备
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2" # 会将实际上0, 2显卡设定为可见：即在该脚本中0 -> 0, 2 -> 1

"""
(2) 补充知识 :: 常见指定device方式

"""
import torch
print(torch.cuda.device_count())

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
a = torch.tensor([2, 2]).to(device=device)

"""
(3) 补充知识 :: 常见的, 直接在GPU上创建torch.tensor的方式

但是基本上用不到，因为基本上我们的数据流向是 硬盘 -> 内存 -> CPU(transform) -> 内存 -> GPU
"""
# 在GPU上直接创建
Y = torch.rand(2, 3, device=device)

x.2 指定 Input Tensor 在哪个设备上

x.2.1 查看 Input Tensor 在哪

print(X.device)

# out >>> device(type='cuda', index=3)

x.2.2 将 Input Tensor 转移到特定设备

X = X.to(device); y = y.to(device)

x.3 指定 Model 在哪个设备上

x.3.1 查看 Model 在哪

通过模型的parameters参数来查看模型在哪。

print(next(self.net.parameters()).device)

x.3.2 将 Model 转移到特定设备

model = model.to(device)
model.to(devie)	# 这两种方式都是可以的，不推荐