目录
一、Pytorch中怎么调用cuda?
1. 第一种
if torch.cuda.is_available():
my_module = my_module.cuda()
2. 第二种
# 定义设备
device = torch.device("cuda")
or
device = torch.device("cuda:0") # 选择具体哪个gpu,默认为0
# 使用
my_module = my_module.to(device)
二、哪些对象可以使用cuda?
1.网络模型
(module)
device = torch.device("cuda")
module = module.to(device)
2.数据 (输入imgs、标注target/label)
training时开启cuda,则testing/validation时也要开启cuda,否则会报错
device = torch.device("cuda")
# training
for data in data_loader:
imgs, targets = data
imgs = imgs.to(device)
targets = target.to(device)
#....
# testing/validation
# 同上
3.损失函数 (loss_function)
device = torch.device("cuda")
loss_function = loss_function.to(device)