torch.cuda

创建日期：2016年12月23日 | 最后更新日期：2025年6月13日

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

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

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

CUDA 语义 包含有关使用 CUDA 的更多详细信息。

StreamContext	选择给定流的上下文管理器。
can_device_access_peer	检查两个设备之间是否可以进行对等访问。
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	在 CUDA IPC 释放 GPU 内存后，强制收集 GPU 内存。
is_available	返回一个布尔值，指示 CUDA 当前是否可用。
is_initialized	返回 PyTorch 的 CUDA 状态是否已初始化。
is_tf32_supported	返回一个布尔值，指示当前 CUDA/ROCm 设备是否支持 dtype 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 事件的包装器。

图 (beta)

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_CUDA_ALLOC_CONF 设置的活动分配器后端的字符串。
CUDAPluggableAllocator	从 so 文件加载的 CUDA 内存分配器。
change_current_allocator	将当前使用的内存分配器更改为提供的分配器。
MemPool	MemPool 表示缓存分配器中的内存池。

caching_allocator_enable

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

class torch.cuda.use_mem_pool(pool, device=None)

一个上下文管理器，它将分配路由到给定的内存池。

参数

• pool (torch.cuda.MemPool) – 要激活的 MemPool 对象，以便分配路由到此池。

• device (torch.device 或 int, 可选) – 所选设备。如果 device 为 None（默认），则使用由 current_device() 给出的当前设备上的 MemPool。

注意

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

NVIDIA 工具扩展 (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 提供的。某些运算符已使用多种策略作为 Tunable 运算符实现。在运行时，所有策略都会被分析，并选择最快的策略用于所有后续操作。

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

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 的包装器。