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 设备是否支持 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 提供。
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 事件包装器

图（测试版）¶

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 （默认），则使用当前设备的 MemPool，由 current_device() 给出。

NVIDIA 工具扩展（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 内核，该内核支持返回一个或多个输出，用于元素级操作。

可调操作符

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

查看文档以获取如何使用它的信息。

流消毒器（原型）¶

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

GPUDirect 存储（原型）¶

torch.cuda.gds 中的 API 提供了对某些 cuFile API 的薄包装，允许在 GPU 内存和存储之间进行直接内存访问传输，避免了 CPU 中的弹跳缓冲区。有关详细信息，请参阅 cufile API 文档。

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

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

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