Pytorch中调用cuda的方法及可以使用cuda的对象

目录

一、Pytorch中怎么调用cuda？？

1. 第一种

2. 第二种

二、哪些对象可以使用cuda?

1.网络模型 (module)

2.数据 (输入imgs、标注target/label)

3.损失函数 (optimizer)

---

一、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)

---