快捷方式

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-smiamd-smi 给出的已用全局(设备)内存。

device_of

上下文管理器,用于将当前设备更改为给定对象的设备。

get_arch_list

返回此库编译所针对的 CUDA 架构列表。

get_device_capability

获取设备的 CUDA 功能。

get_device_name

获取设备的名称。

get_device_properties

获取设备的属性。

get_gencode_flags

返回此库编译时使用的 NVCC gencode 标志。

get_sync_debug_mode

返回 CUDA 同步操作的调试模式的当前值。

init

初始化 PyTorch 的 CUDA 状态。

ipc_collect

强制收集已被 CUDA IPC 释放的 GPU 内存。

is_available

返回一个布尔值,指示 CUDA 当前是否可用。

is_initialized

返回 PyTorch 的 CUDA 状态是否已初始化。

memory_usage

返回在过去的采样周期内,全局(设备)内存被读取或写入的时间百分比,由 nvidia-smi 给出。

set_device

设置当前设备。

set_stream

设置当前流。这是一个用于设置流的包装器 API。

set_sync_debug_mode

设置 CUDA 同步操作的调试模式。

stream

围绕 Context-manager StreamContext 包装,用于选择给定的流。

synchronize

等待 CUDA 设备上所有流中的所有内核完成。

utilization

返回在过去的采样周期内,一个或多个内核在 GPU 上执行的时间百分比,由 nvidia-smi 给出。

temperature

返回 GPU 传感器的平均温度,单位为摄氏度 (°C)。

power_draw

返回 GPU 传感器的平均功耗,单位为毫瓦 (mW)。

clock_rate

返回 GPU SM 在过去采样周期内的时钟速度,单位为赫兹 (Hz),由 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.Modules)并返回图版本。

内存管理

empty_cache

释放缓存分配器当前持有的所有未占用的缓存内存,以便这些内存可以在其他 GPU 应用程序中使用,并在 nvidia-smi 中可见。

get_per_process_memory_fraction

获取进程的内存比例。

list_gpu_processes

返回给定设备的正在运行的进程及其 GPU 内存使用情况的人类可读打印输出。

mem_get_info

使用 cudaMemGetInfo 返回给定设备的全局可用和总 GPU 内存。

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 内存分配器跟踪的“峰值”统计信息。

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)[源代码][源代码]

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

参数

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 内核。

TunableOp

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

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

流清理器(原型)

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得您的问题解答

查看资源