torch.cuda

这个包提供了对CUDA张量的支持,实现与CPU张量相同的功能,利用GPU进行计算。包括当前CUDA设备、流、内存管理和随机数生成器的相关函数,如current_device(), memory_allocated(), set_rng_state()等,适用于深度学习和机器学习场景。" 125211589,9031025,Elasticsearch 冷热数据分离策略,"['Elasticsearch', '大数据', '搜索引擎']

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

参考 torch.cuda - 云+社区 - 腾讯云

目录

torch.cuda.current_blas_handle()

torch.cuda.current_device()

torch.cuda.current_stream(device=None)

torch.cuda.default_stream(device=None)

class torch.cuda.device(device)

torch.cuda.device_count()

class torch.cuda.device_of(obj)

torch.cuda.empty_cache()

torch.cuda.get_device_capability(device=None)

torch.cuda.get_device_name(device=None)

torch.cuda.init()

torch.cuda.ipc_collect()

torch.cuda.is_available()

torch.cuda.max_memory_allocated(device=None)

torch.cuda.max_memory_cached(device=None)

torch.cuda.memory_allocated(device=None)

torch.cuda.memory_cached(device=None)

torch.cuda.reset_max_memory_allocated(device=None)

torch.cuda.reset_max_memory_cached(device=None)

torch.cuda.set_device(device)

torch.cuda.stream(stream)

torch.cuda.synchronize(device=None)

Random Number Generator

torch.cuda.get_rng_state(device='cuda')

torch.cuda.get_rng_state_all()

torch.cuda.set_rng_state(new_state, device='cuda')

torch.cuda.set_rng_state_all(new_states)

torch.cuda.manual_seed(seed)

torch.cuda.seed()

torch.cuda.seed_all()

torch.cuda.initial_seed()

Communication collectives

torch.cuda.comm.broadcast(tensor, devices)

torch.cuda.comm.broadcast_coalesced(tensors, devices, buffer_size=10485760)

torch.cuda.comm.reduce_add(inputs, destination=None)

torch.cuda.comm.scatter(tensor, devices, chunk_sizes=None, dim=0, streams=None)

torch.cuda.comm.gather(tensors, dim=0, destination=None)

Streams and events

class torch.cuda.Stream

query()

record_event(event=None)

synchronize()

wait_event(event)

wait_stream(stream)

class torch.cuda.Event

elapsed_time(end_event)

ipc_handle()

query()

record(stream=None)

synchronize()

wait(stream=None)

Memory management

torch.cuda.empty_cache()

torch.cuda.memory_allocated(device=None)

torch.cuda.max_memory_allocated(device=None)

torch.cuda.reset_max_memory_allocated(device=None)

torch.cuda.memory_cached(device=None)

torch.cuda.max_memory_cached(device=None)

torch.cuda.reset_max_memory_cached(device=None)

NVIDIA Tools Extension (NVTX)

torch.cuda.nvtx.mark(msg)

torch.cuda.nvtx.range_push(msg)

torch.cuda.nvtx.range_pop()


这个包增加了对CUDA张量类型的支持,它实现了与CPU张量相同的功能,但是它们利用gpu进行计算。它是惰性初始化的,所以你总是可以导入它,并使用is_available()来确定您的系统是否支持CUDA。CUDA语义提供了更多关于使用CUDA的细节。

torch.cuda.current_blas_handle()

返回指向当前cuBLAS句柄的cublasHandle_t指针。

torch.cuda.current_device()

返回当前选定设备的索引。

torch.cuda.current_stream(device=None)

返回给定设备当前选定的流。

参数:

  • device (torch.device or int, optional) – 选定的设备。返回当前设备当前选择的流,如果设备为None(默认),则由current_device()给出。

torch.cuda.default_stream(device=None)

返回给定设备的默认流。

参数:

  • device (torch.device or int, optional) – 选定的设备。返回当前设备的默认流,如果设备为None(默认),则由current_device()提供。

class torch.cuda.device(device)

更改所选设备的上下文管理器。

参数:

  • device (torch.device or int) – 要选择的设备索引。如果这个参数是负整数或None,那么它就是no-op。

torch.cuda.device_count()

返回可用的gpu数量。

class torch.cuda.device_of(obj)

将当前设备更改为给定对象的设备的上下文管理器。您可以同时使用张量和存储作为参数。如果一个给定的对象没有分配在GPU上,这是一个no-op。

参数:

obj (Tensor or Storage) – 在选定设备上分配的对象。

torch.cuda.empty_cache()

释放缓存分配器当前持有的所有未占用的缓存内存,以便这些内存可以在其他GPU应用程序中使用,并在nvidia-smi中可见。

注意

empty_cache()不会增加PyTorch可用的GPU内存。有关GPU内存管理的更多细节,请参见内存管理。

torch.cuda.get_device_capability(device=None)

获取设备的cuda功能。

参数:

  • device (torch.device or int, optional) – 用于返回设备功能的设备。如果这个参数是一个负整数,那么这个函数就是no-op。如果设备为None(默认),则使用current_device()提供的当前设备。

返回值:

  • 主要和次要cuda功能的设备,返回类型。

tuple(int, int)

torch.cuda.get_device_name(device=None)

获取设备的名称。

参数:

  • device (torch.device or int, optional) – 用于返回名称的设备。如果这个参数是一个负整数,那么这个函数就是no-op。如果设备为None(默认),则使用current_device()提供的当前设备。

torch.cuda.init()

初始化PyTorch的CUDA状态。如果您通过PyTorch的C API与它进行交互,可能需要显式地调用这个函数,因为在初始化之前,CUDA功能的Python绑定不会这样做。普通用户不应该需要这样做,因为PyTorch的所有CUDA方法都会根据需要自动初始化CUDA状态。如果CUDA状态已经初始化,则不执行任何操作。

torch.cuda.ipc_collect()

Force在CUDA IPC释放GPU内存后收集GPU内存。

注意:

检查是否有任何已发送的CUDA张量可以从内存中清除。如果没有活动计数器,则强制关闭用于引用计数的共享内存文件。当生成器进程停止主动发送张量并希望释放未使用的内存时,此函数非常有用。

torch.cuda.is_available()

返回一个bool,指示CUDA当前是否可用。

torch.cuda.max_memory_allocated(device=None)

返回给定设备张量占用的最大GPU内存(以字节为单位)。默认情况下,这将返回自该程序开始以来分配的内存峰值。reset_max_memory_assigned()可用于重置跟踪此指标的起始点。例如,这两个函数可以测量训练循环中每个迭代的分配内存使用量峰值。

参数:

  • device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。

注意

有关GPU内存管理的更多细节,请参见内存管理。

torch.cuda.max_memory_cached(device=None)

返回缓存分配器为给定设备管理的最大GPU内存(以字节为单位)。默认情况下,这将返回自该程序开始以来的峰值缓存内存。reset_max_memory_cached()可用于重置跟踪此指标的起始点。例如,这两个函数可以测量训练循环中每个迭代的峰值缓存内存量。

参数:

  • device (to
<think>我们正在讨论的是`torch.cuda.device`,它是一个上下文管理器,用于临时切换默认的CUDA设备。 根据引用[1]和引用[2]中关于设备使用的上下文,我们可以结合这些信息来回答。 注意:引用[3]提到的是CUDA不可用的问题,但我们的问题是关于`torch.cuda.device`的用法,所以主要参考引用[1]和[2]。 功能:`torch.cuda.device` 是一个上下文管理器,用于在代码块内部临时切换默认的CUDA设备。当进入该上下文时,所有在上下文中创建的CUDA张量将被分配到指定的设备上。它不会影响已经存在的张量,也不会改变当前线程的当前CUDA流。 使用方法: 1. 导入torch 2. 使用`with torch.cuda.device(id):`来创建一个上下文,其中`id`是GPU设备的索引(整数)或一个设备对象(如`torch.device('cuda:0')`)。 3. 在该上下文中创建的张量将默认位于指定的设备上。 示例代码: import torch # 假设我们有两块GPU device0 = torch.device('cuda:0') device1 = torch.device('cuda:1') # 默认情况下,当前设备是0(如果没有设置过的话) # 使用上下文管理器临时切换到设备1 with torch.cuda.device(device1): # 在设备1上创建张量 a = torch.tensor([1.0, 2.0], device='cuda') # 注意:这里可以不写索引,因为上下文已经指定了默认设备 # 也可以直接使用设备对象 b = torch.tensor([3.0, 4.0], device=device1) # 效果同上 # 在上下文外,默认设备又恢复为原来的设备(设备0) c = torch.tensor([5.0, 6.0], device='cuda') # 在设备0上创建 # 验证 print(a.device) # 输出:cuda:1 print(b.device) # 输出:cuda:1 print(c.device) # 输出:cuda:0 注意:这个上下文管理器只影响默认的设备选择。如果显式指定了设备(如`device='cuda:0'`),则优先使用显式指定的设备。 另外,我们也可以使用整数作为设备索引: with torch.cuda.device(1): d = torch.tensor([7.0, 8.0], device='cuda') # 在设备1上 重要:这个上下文管理器不会改变当前线程的当前CUDA流(关于CUDA流,请参考PyTorch文档)。 另外,我们还可以使用`torch.cuda.set_device(device_index)`来设置默认的CUDA设备,但它是全局设置,会影响整个线程。而`torch.cuda.device`上下文管理器是局部的,只影响其内部的代码块。 引用[1]提到:通过不断学习和实践,你将能够更有效地利用GPU资源来加速你的深度学习任务。因此,合理使用设备上下文管理器可以帮助我们更好地管理多GPU任务。 现在,根据用户的问题,我们回答`torch.cuda.device`的作用及使用方法。 同时,生成相关问题。</think>### PyTorch 中 `torch.cuda.device` 的作用及使用方法详解 `torch.cuda.device` 是 PyTorch 提供的**设备上下文管理器**,用于在代码块中临时切换默认的 CUDA 设备。其核心功能是**动态控制 GPU 计算资源的分配**,特别适用于多 GPU 环境下的设备管理。 --- #### 一、核心功能 1. **临时设备切换** 在 `with` 语句块内创建的所有 CUDA 张量/模型会自动分配到指定设备 ```python # 在设备1上创建张量 with torch.cuda.device(1): tensor1 = torch.randn(3) # 自动创建在 GPU 1 ``` 2. **避免显式指定设备** 简化代码,避免重复书写 `device='cuda:X'` ```python # 等效于显式指定 tensor2 = torch.randn(3, device=torch.device('cuda:1')) ``` 3. **嵌套设备控制** 支持多层上下文管理,内层覆盖外层设置 ```python with torch.cuda.device(0): # 外层:设备0 with torch.cuda.device(1): # 内层:设备1 model = MyModel() # 创建在 GPU 1 data = torch.randn(10, 10) # 创建在 GPU 0 ``` --- #### 二、使用方法 ##### 基础语法 ```python import torch with torch.cuda.device(device_id): # device_id 为整数或设备对象 # 在此区域内创建的 CUDA 对象自动分配到指定设备 ... ``` ##### 典型应用场景 1. **单 GPU 环境指定设备** ```python # 确保 CUDA 可用 [^2] if torch.cuda.is_available(): with torch.cuda.device(0): # 指定 GPU 0 model = MyModel() # 模型在 GPU 0 data = data.cuda() # 数据迁移到 GPU 0 ``` 2. **多 GPU 并行计算** ```python # 在 GPU 0 和 GPU 1 分别执行任务 results = [] for i in range(2): with torch.cuda.device(i): x = torch.randn(1000, device='cuda') # 自动选当前设备 results.append(x * 2) ``` 3. **与 `torch.cuda.Event` 配合计时** [^2] ```python start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) with torch.cuda.device(1): start_event.record() # GPU 1 上的计算操作 result = complex_operation() end_event.record() torch.cuda.synchronize() print(f"耗时: {start_event.elapsed_time(end_event)} ms") ``` --- #### 三、注意事项 1. **作用域限制** 只影响 `with` 块内**新创建**的对象,不影响已存在的对象 ```python tensor_cpu = torch.randn(3) with torch.cuda.device(1): tensor_gpu = tensor_cpu.cuda() # tensor_gpu 在 GPU 1 # tensor_cpu 仍在 CPU ``` 2. **设备索引规则** - 有效范围:`0` 到 `torch.cuda.device_count()-1` - 非法索引会引发 `RuntimeError` 3. **与全局设置的优先级** 上下文管理器优先级高于 `torch.cuda.set_device()` ```python torch.cuda.set_device(0) # 全局默认设备0 with torch.cuda.device(1): # 此处使用设备1(覆盖全局设置) ``` 4. **CPU/GPU 兼容写法** ```python device_id = 0 if torch.cuda.is_available() else None with torch.cuda.device(device_id) if device_id is not None else nullcontext(): # 自动适配环境 ``` --- #### 四、常见问题解决 当遇到 `torch.cuda.is_available() == False` 时 [^3]: 1. 检查 CUDA 与 PyTorch 版本匹配 2. 验证 NVIDIA 驱动兼容性 3. 重装 `cudatoolkit`: ```bash conda install cudatoolkit=11.3 -c nvidia ``` > **最佳实践**:结合 `torch.cuda.device` 和 `torch.device` 对象,可实现灵活的设备管理策略[^1]。 ---
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

Wanderer001

ROIAlign原理

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值