给网络模型、损失函数、数据(输入、标注)加上cuda(注意要赋值给原变量):
法一:
直接加.cuda()
if torch.cuda.is_available(): 防止因为没有GPU而报错
#法一:网络模型、数据(输入、标注)、损失函数 +.cuda()
xigua1=Xigua()
if torch.cuda.is_available():xigua1=xigua1.cuda() #网络模型+cudaloss_func=nn.CrossEntropyLoss()
if torch.cuda.is_available():loss_func=loss_func.cuda() #损失函数+cudafor data in train_dataloader:imgs,targets=dataif torch.cuda.is_available():imgs=imgs.cuda() #数据+cudatargets=targets.cuda()
法二:
先定义device=torch.device("cuda")
再加.to(device)
参数可以写"cpu" "cuda" "cuda:0" "cuda:1"或
"cuda" if torch.cuda.is_available() else "cpu" 防止因为没有GPU而报错
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")xigua1=Xigua()
xigua1=xigua1.to(device) #把网络模型转移到device设备上loss_func=nn.CrossEntropyLoss()
loss_func=loss_func.to(device) #把损失函数转移到device设备上for data in train_dataloader:imgs,targets=data imgs=imgs.to(device) #把数据转移到device设备上targets=targets.to(device)
免费使用GPU的网站:
https://colab.research.google.com/
在Edit的Notebook settings里可以修改硬件配置,选择GPU