torch.cuda

此包增加了对 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 设备是否支持 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 传感器以摄氏度 (Degrees C) 为单位的平均温度。
power_draw	返回 GPU 传感器以毫瓦 (MilliWatts) 为单位的平均功耗。
clock_rate	返回在过去采样周期内，GPU SM 以兆赫兹 (megahertz) 为单位的时钟速度，由 nvidia-smi 报告。
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.scatter	将张量分散到多个 GPU。
comm.gather	从多个 GPU 设备收集张量。

流和事件

Stream	CUDA 流的包装器。
ExternalStream	外部分配的 CUDA 流的包装器。
Event	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）并返回图化版本。

内存管理

empty_cache	释放当前由缓存分配器持有的所有未占用的缓存内存，以便这些内存可用于其他 GPU 应用并在 nvidia-smi 中可见。
get_per_process_memory_fraction	获取进程的内存比例。
list_gpu_processes	返回给定设备的运行进程及其 GPU 内存使用的可读性报告。
mem_get_info	使用 cudaMemGetInfo 返回给定设备的全局空闲和总 GPU 内存。
memory_stats	返回给定设备的 CUDA 内存分配器统计信息字典。
host_memory_stats	返回给定设备的 CUDA 内存分配器统计信息字典。
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 表示缓存分配器中的一个内存池。
MemPoolContext	MemPoolContext 保存当前活动的内存池并暂存之前的内存池。

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。

NVIDIA 工具扩展 (NVTX)

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

Jiterator (beta)

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

TunableOp

某些操作可以使用不止一个库或不止一种技术来实现。例如，GEMM 可以针对 CUDA 或 ROCm 使用 cublas/cublasLt 库或 hipblas/hipblasLt 库来实现。如何知道哪种实现最快并且应该被选择？这就是 TunableOp 提供的功能。某些算子已作为 Tunable Operators 使用多种策略实现。在运行时，会对所有策略进行性能分析，并选择最快的一种用于所有后续操作。

请参阅 文档 以获取如何使用它的信息。

流检查器 (原型)

CUDA 检查器是一个用于检测 PyTorch 中流之间同步错误的原型工具。请参阅 文档 以获取如何使用它的信息。

GPUDirect Storage (原型)

位于 torch.cuda.gds 中的 API 提供了某些 cuFile API 的轻量级包装器，允许在 GPU 内存和存储之间进行直接内存访问传输，避免在 CPU 中使用跳板缓冲区 (bounce buffer)。请参阅 cufile API 文档 以获取更多详细信息。

这些 API 可用于 CUDA 12.6 及更高版本。为了使用这些 API，必须确保系统已按照 GPUDirect Storage 文档 的说明进行了适当配置以使用 GPUDirect Storage。

请参阅 GdsFile 的文档以获取如何使用这些 API 的示例。

gds_register_buffer	在 CUDA 设备上将存储注册为 cufile 缓冲区。
gds_deregister_buffer	在 CUDA 设备上注销先前已注册的存储作为 cufile 缓冲区。
GdsFile	cuFile 的包装器。