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)