需要在Pytorch下跑一SalGAN,但是代码居然用CPU来计算,将其改为GPU版。Pytorch下使用GPU计算需要三步,记录一下:
1.将数据载入GPU:
input -> input.cuda()
2.将模型载入GPU
model = model.cuda()
3.将输出的预测转为numpy格式
result.data.numpy() -> result.data.cpu().numpy()
最后
以上就是留胡子花瓣最近收集整理的关于Pytorch CPU运行与GPU运行的全部内容,更多相关Pytorch内容请搜索靠谱客的其他文章。
本图文内容来源于网友提供,作为学习参考使用,或来自网络收集整理,版权属于原作者所有。
发表评论 取消回复