您当前的位置:首页 > 计算机 > 编程开发 > Python

将gpu运行的torch程序改为cpu运行

时间:03-18来源:作者:点击数:

拿到了师兄的程序,但是自己手头的电脑无法使用cpu训练,虽然也可以使用colab在线运行,但是还是本地的会舒服一些,anyway,就得将代码稍作一下修改

方法:

在训练之前加入这一句

#选择使用gpu,如果没有gpu就使用cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

然后,在训练的代码中会找到.cuda(),将所有的都替换成为.to(device)即可

(在网络后边和输入输出后边都需要改)

即将

model=model.cuda()
x=x.cuda()
y=y.cuda()

上边的改为下边的!!!!!!!

model=model.to(device)
x=x.to(device)
y=y.to(device)

如果电脑有gpu可用,调用的还是gpu

方便获取更多学习、工作、生活信息请关注本站微信公众号城东书院 微信服务号城东书院 微信订阅号
推荐内容
相关内容
栏目更新
栏目热门