Allocator【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge产品支持情况产品是否支持Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√模块导入from ge.allocator import Allocator, MemBlock功能说明Allocator 是外部内存分配器的抽象基类用于自定义设备内存管理策略。MemBlock 表示已分配的内存块持有设备内存地址和大小。用户需要继承 Allocator 类并实现 malloc 和 free 方法然后通过 Session.register_external_allocator() 注册到指定 stream。类定义MemBlock 类class MemBlock: def __init__(self, addr: int, size: int) property def addr(self) - int property def size(self) - intAllocator 类class Allocator(ABC): abstractmethod def malloc(self, size: int) - MemBlock abstractmethod def free(self, block: MemBlock) - None函数说明MemBlock 类__init__def __init__(self, addr: int, size: int)功能说明创建内存块实例持有设备内存地址和大小信息。参数说明参数名类型必选/可选说明addrint必选设备内存地址。sizeint必选内存大小单位为字节。返回值说明无返回值。addr属性property def addr(self) - int功能说明获取设备内存地址。参数说明无参数。返回值说明返回值类型说明int设备内存地址。size属性property def size(self) - int功能说明获取内存大小。参数说明无参数。返回值说明返回值类型说明int内存大小单位为字节。Allocator 类mallocabstractmethod def malloc(self, size: int) - MemBlock功能说明分配指定大小的设备内存返回包含设备内存地址和大小的 MemBlock 对象。此为抽象方法子类必须实现。参数说明参数名类型必选/可选说明sizeint必选待分配的内存大小单位为字节。返回值说明返回值类型说明MemBlock已分配的内存块对象包含有效的设备内存地址。约束说明子类必须实现此方法否则无法实例化。分配失败时应抛出 MemoryError。freeabstractmethod def free(self, block: MemBlock) - None功能说明释放之前通过 malloc 分配的设备内存。此为抽象方法子类必须实现。参数说明参数名类型必选/可选说明blockMemBlock必选待释放的内存块对象应为之前通过 malloc 返回的 MemBlock 实例。返回值说明无返回值。约束说明子类必须实现此方法否则无法实例化。不应对同一内存块重复释放。【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考