利用GPU进行模型训练

GPU训练方式1

网络模型
数据（输入，标注）
损失函数
加上.cuda()

GPU训练方式2

# 测试是否有gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 如果有多张显卡，可以指定显卡
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
print(device)
.to(device)
posted @ 2022-10-07 14:02  watsono  阅读(58)  评论(0)  编辑  收藏  举报