cpu与gpu的切换

最新推荐文章于 2023-11-25 20:21:26 发布

科研小达人

最新推荐文章于 2023-11-25 20:21:26 发布

阅读量2.3k

点赞数

分类专栏： python 文章标签：深度学习机器学习人工智能

本文链接：https://blog.youkuaiyun.com/weixin_43231756/article/details/125771553

版权

python 专栏收录该内容

9 篇文章

订阅专栏

该博客介绍了如何在PyTorch中选择设备进行训练，特别是在有GPU的情况下，通过设置`device`变量来自动在CPU和GPU之间切换。文章详细讲解了如何将模型、输入数据和目标标签转换到选定的设备上，以实现更高效的计算。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

方法：
在训练之前加入这一句

#选择使用gpu，如果没有gpu就使用cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
1
2
然后，在训练的代码中会找到 .cuda() ，将所有的都替换成为.to(device) 即可
（在网络后边和输入输出后边都需要改）

即将

model=model.cuda()
x=x.cuda()
y=y.cuda()
1
2
3
上边的改为下边的！！！！！！！

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