利用GPU训练

目录

利用GPU训练，只需要对三个部分转到cuda即可

方法一：.cuda() 

方法二：.to(device)

---

利用GPU训练，只需要对三个部分转到cuda即可

1. 模型

2. Loss函数

3. imgs 和targets，这部分需要在训练和验证时进行转移

方法一：.cuda() 

# 创建网络模型
module = Module()
# 网络模型转移到cuda
module = module.cuda()
# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 损失函数放在cuda
loss_fn = loss_fn.cuda()

    for data in train_data_Loader:
        imgs ,targets = data
        # 训练数据转cuda
        imgs = imgs.cuda()
        targets = targets.cuda()
        output = module(imgs)

 for data in test_data_Loader:
            imgs ,targets = data
            imgs = imgs.cuda()
            targets = targets.cuda()
            output = module(imgs)

方法二：.to(device)

if torch.cuda.is_available():
    module = module.cuda()
#或者

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

• 如果有多个显卡，可以先定义GPU

device = torch.device('cuda:0')

device = torch.device('cuda:1')

# 创建网络模型
module = Module()
# 网络模型转移到cuda
module = module.to(deivce)
# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 损失函数放在cuda
loss_fn = loss_fn.to(device)