评价此页

torch.mps#

创建于:2023 年 2 月 10 日 | 最后更新于:2025 年 6 月 8 日

此包提供了一个在 Python 中访问 MPS(Metal Performance Shaders)后端接口的接口。Metal 是 Apple 的编程 Metal GPU(图形处理器单元)的 API。使用 MPS 意味着可以通过在 Metal GPU 上运行工作来提高性能。有关更多详细信息,请参阅 https://developer.apple.com/documentation/metalperformanceshaders

device_count

返回可用 MPS 设备的数量。

synchronize

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

get_rng_state

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

set_rng_state

设置随机数生成器状态。

manual_seed

设置用于生成随机数的种子。

seed

将生成随机数的种子设置为一个随机数。

empty_cache

释放缓存分配器当前持有的所有未占用的缓存内存,以便其他 GPU 应用程序可以使用它们。

set_per_process_memory_fraction

设置内存分数以限制进程在 MPS 设备上的内存分配。

current_allocated_memory

以字节为单位返回当前被张量占用的 GPU 内存。

driver_allocated_memory

返回 Metal 驱动程序为该进程分配的总 GPU 内存(以字节为单位)。

recommended_max_memory

返回 GPU 内存的推荐最大工作集大小(以字节为单位)。

compile_shader

从源文件编译计算着色器,并允许从 Python 运行时方便地调用其中定义的内核。示例。

MPS 性能分析器#

profiler.start

从 MPS 后端开始 OS Signpost 跟踪。

profiler.stop

停止从 MPS 后端生成 OS Signpost 跟踪。

profiler.profile

上下文管理器,用于启用从 MPS 后端生成 OS Signpost 跟踪。

profiler.is_capturing_metal

检查 Metal 捕获是否正在进行中

profiler.is_metal_capture_enabled

检查 metal_capture 上下文管理器是否可用。要启用 Metal 捕获,请设置 MTL_CAPTURE_ENABLED 环境变量。

profiler.metal_capture

上下文管理器,用于将 Metal 调用捕获到 gputrace 中

MPS 事件#

event.Event

MPS 事件的包装器。