Pytorch如何指定device(cuda or cpu)
更新时间：2024年06月27日 16:59:07 作者：湫兮之风
这篇文章主要介绍了Pytorch如何指定device(cuda or cpu)问题,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教
Pytorch指定device(cuda or cpu)
torch.device()
在PyTorch中，你可以使用torch.device函数来设置设备（device）。
这个函数接受一个字符串作为参数，该字符串可以是 "cpu" 或 "cuda"，它指定了要在哪个设备上运行张量计算。
# 使用CPU
device = torch.device("cpu")
# 使用GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cuda:0")
device = torch.device("cuda:1")os.environ['CUDA_VISIBLE_DEVICES']
全局指定只能识别的cuda编号
os.environ['CUDA_VISIBLE_DEVICES']
是一个在运行 Python 程序时可以设置可见的 CUDA 设备的环境变量。
这个环境变量可以控制哪些 GPU 设备对 Python 程序是可见的。
例如：
如果你想让你的程序只能看到编号为 0 的 GPU，你可以在运行程序前设置这个环境变量：
import os os.environ['CUDA_VISIBLE_DEVICES'] = '0'
或者你也可以设置多个 GPU，只需将它们的编号以逗号分隔：
import os os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
如果你想让程序对所有 GPU 都可见，只需将环境变量设置为空字符串：
os.environ['CUDA_VISIBLE_DEVICES'] = ''
说说torch.device的作用
torch.device代表将torch.Tensor分配到的设备的对象。
torch.device包含一个设备类型（‘cpu’或‘cuda’）和可选的设备序号。
如果设备序号不存在，则为当前设备。
如：torch.Tensor用设备构建‘cuda’的结果等同于‘cuda：X’，其中X是torch.cuda.current_device()的结果。
#开始脚本，创建一个张量
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
...
#但是无论你获得一个新的Tensor或者Module
#如果他们已经在目标设备上则不会执行复制操作
input = data.to(device)
model = MyModule(...).to(device)总结
以上为个人经验，希望能给大家一个参考，也希望大家多多支持脚本之家。
相关文章
这篇文章主要介绍了Python集合的基础操作，Python中的集合和数学上的集合基本是没有区别的，是无序的，即不可以使用索引访问的，集合中是不能出现重复元素的。想着情了解具体内容的小伙伴可以参考下面文章内容2021-11-11
这篇文章主要介绍了python统计日志ip访问数的方法,涉及Python操作日志文件及正则匹配的相关技巧,非常具有实用价值,需要的朋友可以参考下2015-07-07
今天小编就为大家分享一篇TensorFlow2.0:张量的合并与分割实例，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧2020-01-01
这篇文章主要为大家介绍了python深度学习tensorflow实例数据下载与读取示例详解，有需要的朋友可以借鉴参考下，希望能够有所帮助，祝大家多多进步，早日升职加薪2022-06-06
这篇文章主要为大家详细介绍了python+opencv实现摄像头调用的方法，具有一定的参考价值，感兴趣的小伙伴们可以参考一下2019-06-06
大家好，本篇文章主要讲的是python批量修改xml文件中的信息，感兴趣的同学赶快来看一看吧，对你有帮助的话记得收藏一下的相关资料2022-01-01
本文主要介绍了基于OpenCV(python)的实现文本分割之垂直投影法，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧2022-08-08
代码需要经历开发、调试、审查、测试或者上线等不同阶段，在“测试”时，可能只想看警告和错误信息，然而在“调试”时，可能还想看到跟调试相关的信息。如果想打印出使用的模块以及代码运行的时间，那么代码很容易变得混乱。使用logging日志模块，就能很容易地解决2021-06-06
这篇文章主要介绍了关于脚本测试postman快速导出python接口测试示例的过程操作，有需要的朋友可以借鉴参考下，希望能够有所帮助2021-09-09
Python EasyDict库以属性方式访问字典元素(无需使用方括号和键)
在Python中,字典（dict）是一种常用的数据结构,用于存储键值对,然而,有时候我们希望以属性的方式访问字典中的元素,而无需使用方括号和键,这就是EasyDict库的用武之地,本文将深入介绍EasyDict库,展示其强大的功能和如何通过示例代码更好地利用它2023-12-12
最新评论