评价此页

torch.mtia#

创建日期:2023年7月11日 | 最后更新日期:2025年6月8日

MTIA 后端在树外实现,此处仅定义接口。

此包提供了一个用于在 python 中访问 MTIA 后端的接口

StreamContext

选择给定流的上下文管理器。

current_device

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

current_stream

返回给定设备当前选择的Stream

default_stream

返回给定设备的默认Stream

device_count

返回可用 MTIA 设备的数量。

init

is_available

如果 MTIA 设备可用则返回 true

is_initialized

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

memory_stats

返回给定设备的 MTIA 内存分配器统计信息的字典。

get_device_capability

以 (主版本, 次版本) 元组形式返回给定设备的功能。

empty_cache

清空 MTIA 设备缓存。

record_memory_history

在 MTIA 分配器上启用/禁用内存分析器

snapshot

返回 MTIA 内存分配器历史的字典

attach_out_of_memory_observer

向 MTIA 内存分配器附加一个内存不足观察器

set_device

设置当前设备。

set_stream

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

stream

包装选择给定流的上下文管理器 StreamContext。

synchronize

等待 MTIA 设备上所有流中的所有作业完成。

device

更改选定设备的上下文管理器。

set_rng_state

设置随机数生成器状态。

get_rng_state

以 ByteTensor 形式返回随机数生成器状态。

DeferredMtiaCallError

流和事件#

事件

查询和记录流状态以识别或控制流之间的依赖关系并测量时间。

按先进先出 (FIFO) 顺序异步执行相应任务的有序队列。