评价此页

★ ★ ★ ★ ★

torch.cuda#

创建于: 2016年12月23日 | 最后更新于: 2025年10月22日

此包添加了对 CUDA 张量类型的支持。

它实现了与 CPU 张量相同的功能，但利用 GPU 进行计算。

它是惰性初始化的，因此您始终可以导入它，并使用 is_available() 来确定您的系统是否支持 CUDA。

有关使用 CUDA 的更多详细信息，请参阅 CUDA 语义。

`StreamContext`	选择给定流的上下文管理器。
`can_device_access_peer`	检查两个设备之间的对等访问是否可行。
`check_error`	如果 CUDA 运行时 API 调用结果不成功，则引发错误。
`current_blas_handle`	返回当前 cuBLAS 句柄的 cublasHandle_t 指针
`current_device`	返回当前选定设备的索引。
`current_stream`	返回给定设备上当前选定的 `Stream`。
`cudart`	检索 CUDA 运行时 API 模块。
`default_stream`	返回给定设备的默认 `Stream`。
`device`	更改选定设备的上下文管理器。
`device_count`	返回可用 GPU 的数量。
`device_memory_used`	返回全局（设备）内存的已使用量（以字节为单位），如 nvidia-smi 或 amd-smi 所示。
`device_of`	将当前设备更改为给定对象设备的上下文管理器。
`get_arch_list`	返回编译该库时使用的 CUDA 架构列表。
`get_device_capability`	获取设备的 CUDA 能力。
`get_device_name`	获取设备的名称。
`get_device_properties`	获取设备的属性。
`get_gencode_flags`	返回编译该库时使用的 NVCC gencode 标志。
`get_stream_from_external`	从外部分配的 CUDA 流返回一个 `Stream`。
`get_sync_debug_mode`	返回当前 CUDA 同步操作的调试模式值。
`init`	初始化 PyTorch 的 CUDA 状态。
`ipc_collect`	在 GPU 内存被 CUDA IPC 释放后强制收集。
`is_available`	返回一个布尔值，指示 CUDA 当前是否可用。
`is_bf16_supported`	返回一个布尔值，指示当前的 CUDA/ROCm 设备是否支持 bfloat16 数据类型。
`is_initialized`	返回 PyTorch 的 CUDA 状态是否已初始化。
`is_tf32_supported`	返回一个布尔值，指示当前的 CUDA/ROCm 设备是否支持 tf32 数据类型。
`memory_usage`	返回过去采样周期内全局（设备）内存被读取或写入的百分比，如 nvidia-smi 所示。
`set_device`	设置当前设备。
`set_stream`	设置当前流。这是一个用于设置流的包装 API。
`set_sync_debug_mode`	设置 CUDA 同步操作的调试模式。
`stream`	包装选择给定流的上下文管理器 StreamContext。
`synchronize`	等待 CUDA 设备上的所有流中的所有内核完成。
`utilization`	返回过去采样周期内一个或多个内核在 GPU 上执行的百分比，如 nvidia-smi 所示。
`temperature`	返回 GPU 传感器平均温度（摄氏度）。
`power_draw`	返回 GPU 传感器平均功耗（毫瓦）。
`clock_rate`	返回过去采样周期内 GPU SM 的时钟速度（兆赫兹），如 nvidia-smi 所示。
`AcceleratorError`	在设备上执行时引发异常
`OutOfMemoryError`	设备内存不足时引发异常

随机数生成器#

`get_rng_state`	将指定 GPU 的随机数生成器状态作为 ByteTensor 返回。
`get_rng_state_all`	返回一个 ByteTensor 列表，表示所有设备的随机数状态。
`set_rng_state`	设置指定 GPU 的随机数生成器状态。
`set_rng_state_all`	设置所有设备的随机数生成器状态。
`manual_seed`	设置当前 GPU 随机数生成的种子。
`manual_seed_all`	设置所有 GPU 随机数生成的种子。
`seed`	将当前 GPU 随机数生成的种子设置为一个随机数。
`seed_all`	将所有 GPU 随机数生成的种子设置为一个随机数。
`initial_seed`	返回当前 GPU 的当前随机种子。

通信集合#

`comm.broadcast`	将张量广播到指定的 GPU 设备。
`comm.broadcast_coalesced`	将一系列张量广播到指定的 GPU。
`comm.reduce_add`	从多个 GPU 求和张量。
`comm.reduce_add_coalesced`	从多个 GPU 求和张量。
`comm.scatter`	将张量分散到多个 GPU。
`comm.gather`	从多个 GPU 设备收集张量。

流和事件#

流

CUDA 流的包装器。

ExternalStream

包装器，用于包装外部分配的 CUDA 流。

事件

包装器，用于包装 CUDA 事件。

图（测试版）#

`is_current_stream_capturing`	如果当前 CUDA 流正在进行 CUDA 图捕获，则返回 True，否则返回 False。
`graph_pool_handle`	返回一个代表图内存池 ID 的不透明令牌。
`CUDAGraph`	CUDA 图的包装器。
`graph`	上下文管理器，将 CUDA 工作捕获到 `torch.cuda.CUDAGraph` 对象中，以便稍后重放。
`make_graphed_callables`	接受可调用对象（函数或 `nn.Module`）并返回图形式的版本。

此包添加了对 CUDA 中实现的设备内存管理的支持。

内存管理#

`empty_cache`	释放缓存分配器当前持有的所有未占用缓存内存，以便它们可供其他 GPU 应用程序使用并在 nvidia-smi 中可见。
`get_per_process_memory_fraction`	获取进程的内存分数。
`list_gpu_processes`	返回给定设备的正在运行的进程及其 GPU 内存使用情况的可读打印输出。
`mem_get_info`	使用 cudaMemGetInfo 返回给定设备的全局空闲和总 GPU 内存。
`memory_stats`	返回给定设备的 CUDA 内存分配器统计信息的字典。
`memory_stats_as_nested_dict`	以嵌套字典的形式返回 `memory_stats()` 的结果。
`reset_accumulated_memory_stats`	重置 CUDA 内存分配器跟踪的“累积”（历史）统计信息。
`host_memory_stats`	返回给定设备的 CUDA 内存分配器统计信息的字典。
`host_memory_stats_as_nested_dict`	以嵌套字典的形式返回 `host_memory_stats()` 的结果。
`reset_accumulated_host_memory_stats`	重置主机内存分配器跟踪的“累积”（历史）统计信息。
`memory_summary`	返回给定设备的当前内存分配器统计信息的可读打印输出。
`memory_snapshot`	返回跨所有设备的 CUDA 内存分配器状态的快照。
`memory_allocated`	返回给定设备上张量当前占用的 GPU 内存（以字节为单位）。
`max_memory_allocated`	返回给定设备上张量占用的最大 GPU 内存（以字节为单位）。
`reset_max_memory_allocated`	重置跟踪给定设备上张量占用的最大 GPU 内存的起点。
`memory_reserved`	返回给定设备上缓存分配器管理的当前 GPU 内存（以字节为单位）。
`max_memory_reserved`	返回给定设备上缓存分配器管理的最大 GPU 内存（以字节为单位）。
`set_per_process_memory_fraction`	设置进程的内存分数。
`memory_cached`	已弃用；请参阅 `memory_reserved()`。
`max_memory_cached`	已弃用；请参阅 `max_memory_reserved()`。
`reset_max_memory_cached`	重置跟踪给定设备上缓存分配器管理的 GPU 最大内存的起点。
`reset_peak_memory_stats`	重置 CUDA 内存分配器跟踪的“峰值”统计信息。
`reset_peak_host_memory_stats`	重置主机内存分配器跟踪的“峰值”统计信息。
`caching_allocator_alloc`	使用 CUDA 内存分配器执行内存分配。
`caching_allocator_delete`	删除使用 CUDA 内存分配器分配的内存。
`get_allocator_backend`	返回一个描述活动分配器后端的字符串，该后端由 `PYTORCH_ALLOC_CONF` 设置。
`CUDAPluggableAllocator`	从 so 文件加载的 CUDA 内存分配器。
`change_current_allocator`	将当前使用的内存分配器更改为提供的分配器。
`MemPool`	MemPool 代表缓存分配器中的内存池。

caching_allocator_enable

启用或禁用 CUDA 内存分配器。

class torch.cuda.use_mem_pool(pool, device=None)[source]#

一个上下文管理器，用于将分配路由到指定的池。

参数:

pool (torch.cuda.MemPool) – 一个 MemPool 对象，将被激活，以便分配路由到此池。
device (torch.device 或 int, optional) – 指定设备。如果 device 为 None（默认），则使用 current_device() 给出的当前设备上的 MemPool。

注意

此上下文管理器仅使当前线程的分配路由到指定的池。如果在上下文管理器内生成新线程（例如，通过调用 backward），则该线程中的分配将不会路由到指定的池。

NVIDIA Tools Extension (NVTX)#

`nvtx.mark`	描述在某个时间点发生的瞬时事件。
`nvtx.range_push`	将范围推送到嵌套范围跨度的堆栈。
`nvtx.range_pop`	从嵌套范围跨度堆栈中弹出范围。
`nvtx.range`	上下文管理器/装饰器，在作用域开始时推入 NVTX 范围，并在作用域结束时弹出。

Jiterator (测试版)#

`jiterator._create_jit_fn`	为元素级操作创建 jiterator 生成的 CUDA 内核。
`jiterator._create_multi_output_jit_fn`	为支持返回一个或多个输出的元素级操作创建 jiterator 生成的 CUDA 内核。

TunableOp#

某些操作可以使用多个库或多种技术来实现。例如，GEMM 可以通过 cublas/cublasLt 库或 hipblas/hipblasLt 库分别实现 CUDA 或 ROCm。如何知道哪种实现最快并应被选择？这就是 TunableOp 提供的。某些运算符已通过多种策略实现为可调运算符。在运行时，所有策略都会被分析，最快的一种会被选定用于所有后续操作。

有关如何使用它的信息，请参阅文档。

Stream Sanitizer (原型)#

CUDA Sanitizer 是一个用于检测 PyTorch 中流之间同步错误的工具原型。有关如何使用它的信息，请参阅文档。

GPUDirect Storage (原型)#

torch.cuda.gds 中的 API 是对某些 cuFile API 的薄封装，它允许 GPU 内存和存储之间进行直接内存访问传输，从而避免了 CPU 中的反弹缓冲区。有关更多详细信息，请参阅 cufile api 文档。

这些 API 可用于 CUDA 版本 12.6 及更高版本。为了使用这些 API，必须确保您的系统已根据 GPUDirect Storage 文档正确配置为使用 GPUDirect Storage。

有关如何使用这些 API 的示例，请参阅 GdsFile 的文档。

gds_register_buffer

将 CUDA 设备上的存储注册为 cufile 缓冲区。

gds_deregister_buffer

取消注册先前在 CUDA 设备上注册为 cufile 缓冲区的存储。

GdsFile

cuFile 的封装。

Green Contexts (实验性)#

torch.cuda.green_contexts 是对 CUDA Green Context API 的薄封装，用于实现更通用的 SM 资源划分给 CUDA 内核。

这些 API 可在 PyTorch 中与 CUDA 版本 12.8 及更高版本一起使用。

有关如何使用这些 API 的示例，请参阅 GreenContext 的文档。

GreenContext

CUDA 绿色上下文的封装。