Pytorch在GPU上训练

Pytorch在指定单GPU上训练

device=torch.device("cuda:0"if torch.cuda.is_available() else "cpu")

确定用cuda：0这块GPU

def Model(...):
	....
	....

model=Model().to(device)
#model=Model().cuda()

模型需要cuda化

inputs, labels= inputs.to(device), labels.to(device)

训练和测试集都必须cuda化

import datetime

starttime = datetime.datetime.now()

#long running

endtime = datetime.datetime.now()

print (endtime - starttime).seconds

通过计算运行时间可以得出 gpu上基本可以加快三倍以上