to(device)

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

将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去,之后的运算都在GPU上进行。
多个GPU: 

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")   model = Model()   if torch.cuda.device_count() > 1:       model = nn.DataParallel(model,device_ids=[0,1,2])          model.to(device)

to(device)最先出现在Python成神之路

版权声明:
作者:Zad
链接:https://www.techfm.club/p/19334.html
来源:TechFM
文章版权归作者所有,未经允许请勿转载。

THE END
分享
二维码
< <上一篇
下一篇>>