目录

torch.cuda

此包增加了对 CUDA 张量类型的支持,这些类型实现了相同的 函数作为 CPU 张量,但它们使用 GPU 进行计算。

它是延迟初始化的,因此您始终可以导入它,并用于确定您的系统是否支持 CUDA。

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

StreamContext (流上下文)

选择给定流的 context-manager。

can_device_access_peer

检查两个设备之间是否可以进行对等访问。

current_blas_handle

返回指向当前 cuBLAS 句柄cublasHandle_t指针

current_device

返回当前所选设备的索引。

current_stream

返回当前为给定设备选择的 ID。

default_stream

返回给定设备的默认值

装置

更改所选设备的 context-manager 。

device_count

返回可用的 GPU 数量。

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 同步操作的调试模式的当前值。

初始化

初始化 PyTorch 的 CUDA 状态。

ipc_collect

Force 在 CUDA IPC 释放 GPU 内存后收集 GPU 内存。

is_available

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

is_initialized

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

memory_usage

返回上一个采样周期中读取或写入全局(设备)内存的时间百分比。

set_device

设置当前设备。

set_stream

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

set_sync_debug_mode

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

围绕选择给定流的 Context-manager StreamContext 的包装器。

同步

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

利用

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

随机数生成器

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 上生成随机数的种子。

种子

将用于生成随机数的种子设置为当前 GPU 的随机数。

seed_all

将用于生成随机数的种子设置为所有 GPU 上的随机数。

initial_seed

返回当前 GPU 的当前随机种子。

通信集体

通信广播

将张量广播到指定的 GPU 设备。

comm.broadcast_coalesced

将序列张量广播到指定的 GPU。

comm.reduce_add

对来自多个 GPU 的张量求和。

comm.scatter

将张量分散到多个 GPU 上。

comm.gather

从多个 GPU 设备收集张量。

流和事件

CUDA 流的包装器。

外部流

包装外部分配的 CUDA 流。

事件

CUDA 事件的包装器。

图形(测试版)

graph_pool_handle

返回表示图形内存池 id 的不透明令牌。

CUDAGraph

CUDA 图形的包装器。

将 CUDA 工作捕获到对象中以供以后重放的上下文管理器。

make_graphed_callables

接受可调用对象 (函数或 s) 并返回图形版本。

内存管理

empty_cache

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

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

荒废的;见

max_memory_cached

荒废的;见

reset_max_memory_cached

重置跟踪给定设备的缓存分配器管理的最大 GPU 内存的起点。

reset_peak_memory_stats

重置 CUDA 内存分配器跟踪的 “峰值” 统计数据。

caching_allocator_alloc

使用 CUDA 内存分配器执行内存分配。

caching_allocator_delete

删除使用 CUDA 内存分配器分配的内存。

NVIDIA 工具扩展 (NVTX)

nvtx.mark

描述在某个时间点发生的瞬时事件。

nvtx.range_push

将范围推送到嵌套范围 span 的堆栈上。

nvtx.range_pop

从嵌套范围 span 堆栈中弹出一个范围。

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

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

查看资源