Pytorch框架学习记录13——利用GPU训练

使用GPU训练的两种方法

• 方法1

数据（输入、标注） 损失函数 模型 .cuda()

# 创建网络模型
test = Test()
test = test.cuda()

# 定义损失函数
loss_fn = torch.nn.CrossEntropyLoss()
loss_fn = loss_fn.cuda()

imgs, target = train_data
imgs = imgs.cuda()
target = target.cuda()

imgs, target = test_data
imgs = imgs.cuda()
target = target.cuda()
output = test(imgs)

• 方法2

.to(device)

Device = torch.device("cpu") 或 torch.device("cuda")

torch.device("cuda:0") 指定第一张显卡

# 定义训练的设备
device = torch.device("cuda:0")

# 创建网络模型
test = Test()
test = test.to(device)

# 定义损失函数
loss_fn = torch.nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)

for train_data in train_loader:
    imgs, target = train_data
    imgs = imgs.to(device)
    target = target.to(device)

for test_data in test_loader:
    imgs, target = test_data
    imgs = imgs.to(device)
    target = target.to(device)