torch.cuda

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

此包增加了对 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 内存。
is_available	返回一个布尔值，指示 CUDA 当前是否可用。
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 的时钟速度（MHz），如“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, optional) – 选定的设备。如果 device 为 None（默认），则使用当前设备上的 MemPool，由 current_device() 提供。

注意

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

NVIDIA Tools Extension (NVTX)

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

Jiterator (beta)

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

TunableOp

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

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

Stream Sanitizer (prototype)

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

GPUDirect Storage (prototype)

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