ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

AMCT 张量分解接口 auto_decomposition 使用指南:PyTorch 模型卷积层自动分解实战

AMCT 张量分解接口 auto_decomposition 使用指南:PyTorch 模型卷积层自动分解实战 AMCT 张量分解接口 auto_decomposition 使用指南PyTorch 模型卷积层自动分解实战【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct导读auto_decomposition是 CANN AMCT昇腾 AI 处理器亲和模型压缩工具仓提供的张量分解接口用于对用户输入的 PyTorch 模型对象执行自动张量分解将符合条件的torch.nn.Conv2d卷积层拆分为多个小卷积层同时完成权重分解并可选地保存分解信息文件供后续复用。本文基于 auto_decomposition 接口文档 与仓库源码完整讲解该接口的产品支持情况、函数原型、参数与返回值、分解约束、调用示例、底层实现原理以及它与配套接口decompose_network的配合用法帮助读者掌握在训练前对模型进行张量分解的完整流程。功能概述对用户输入的 PyTorch 模型对象进行张量分解得到分解后的模型对象和分解前后层的对应名称并保存分解信息文件可选。在 AMCT 的接口体系中张量分解属于独立的模型压缩手段其接口列表定义于 docs/zh/api/README.md接口名称功能描述auto_decomposition对用户输入的 PyTorch 模型对象进行张量分解得到分解后的模型对象和分解前后层的对应名称并保存分解信息文件可选。decompose_network用户输入 PyTorch 模型对象和通过auto_decomposition保存的分解信息文件根据分解信息文件将模型对象改变为张量分解后的结构得到分解后的模型对象和分解前后层的对应名称。两个接口共同构成先分解、后复用分解信息的完整流程auto_decomposition完成分解计算与权重落地decompose_network仅依据分解信息文件重建模型结构权重为torch.nn.Conv2d默认值需另行加载分解后的权重特别适合多进程分布式训练场景。产品支持情况该接口支持以下昇腾产品对应文档见 auto_decomposition.md产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√函数原型model, changes auto_decomposition(model, decompose_info_pathNone)该接口的公共导出路径为amct_pytorch.tensor_decompose实际实现位于amct_pytorch.classic.graph_based.amct_pytorch.tensor_decompose公共包通过重导出暴露文档化的导入路径见 amct_pytorch/tensor_decompose/init.py__all__ [auto_decomposition, decompose_network]参数说明参数名输入/输出使用限制model输入含义待分解的含有预训练权重的 PyTorch 模型对象。在调用该接口时建议将模型放置于 CPU 上以防分解时显存不足。数据类型torch.nn.Module。decompose_info_path输入含义分解信息文件保存路径。将以 JSON 格式存储因此建议使用.json扩展名。为None时不保存分解信息文件默认。数据类型string。默认值None。从源码看tensor_decompose.py_check_auto_decomp_input会对入参做类型校验model必须为torch.nn.Module实例decompose_info_path若传入则必须是字符串并通过os.path.realpath标准化为绝对路径。因此传入非法类型会抛出TypeError如model type is wrong. Expected type torch.nn.Moduledecompose_info_path传None时不落盘分解信息文件仅在打屏日志中记录分解结果。返回值说明张量分解后的模型对象类型为torch.nn.Module。张量分解前后的对应层名构成的字典类型为dict形如{conv1: [conv1.0, conv1.1], conv2: [conv2.0, conv2.1], ...}仅包含发生变化的层。从源码实现看tensor_decompose.py_replace_layers将每个原始层替换为一个nn.Sequential容器内含分解出的first与last两个新卷积层并在日志中打印Decompose: conv1 - [conv1.0, conv1.1]形式的映射关系changes字典即由_log_decomposed_layers逐层构造并合并而来。约束说明用户输入的模型需为torch.nn.Module类型的对象。本接口函数仅对通过torch.nn.Conv2d()构建的卷积进行分解。用户调用本接口函数接口函数对符合分解条件的卷积层进行自动分解约束请参见 分解约束。对应地源码在遍历层时显式执行类型过滤tensor_decompose.pyfor layer_name, layer in model.named_modules(): if not _check_type(layer, nn.Conv2d): continue这意味着分解对象必须是可通过model.named_modules()遍历到的torch.nn.Conv2d模块权重数据类型仅支持torch.float16、torch.float32、torch.float64其余数据类型如整型的卷积层会被跳过并打印 warning 日志Unsupported data type ... ignored分解计算基于 4 维权重张量形状(C_out, C_in, K_h, K_w)若权重张量含inf或nan会抛出ValueError。底层实现原理整体流程auto_decomposition的主流程定义在 tensor_decompose.py大致分为四步入参校验_check_auto_decomp_input校验model与decompose_info_path逐层分解遍历model.named_modules()对每个nn.Conv2d调用_decompose_one_layer返回新层组合new_layer_packs与分解信息decom_info保存分解信息_save_decom_info将分解信息以 JSON 格式缩进为制表符indent\t写入decompose_info_path保存前自动创建父目录权限为rwxr-x---风格无符合条件张量时不生成文件并打印提示替换层结构_replace_layers将原始层替换为nn.Sequential(first_layer, last_layer)同时把新权重迁移到原始权重所在设备返回分解后的模型与changes字典。分解模式判定底层张量分解算法实现在 decomposition.py核心流程如下将权重张量按不同轴展开对展开后的矩阵做SVD奇异值分解获得奇异值序列通过 C 加速库libtensor_decompose.so中的FastFilterConv与GetRank经ctypes调用见 decomposition.py快速过滤不满足分解条件的卷积并依据奇异值计算截断秩rank根据截断结果输出左右两个小权重张量first/last对应DecomposeMode枚举中的五种模式decomposition.py模式含义UNCHANGE不分解不符合分解条件FCFK分解 First Channel 和 First KernelFCSK分解 First Channel 和 Second KernelSCFK分解 Second Channel 和 First KernelSCSK分解 Second Channel 和 Second Kernel若模式为UNCHANGE_decompose_one_layer直接返回空结果该层保持原样。新层构造_create_new_layertensor_decompose.py依据分解结果构造新的torch.nn.Conv2d根据模式确定被修改的维度索引宽或高将对应的stride置 1、padding置 0、dilation置 1另一维保持原值从而保证分解前后卷积输出张量的空间尺寸一致仅last位置的新层携带 biasfirst位置新层 bias 恒为False通过param_dict处理共享权重同一权重被多个卷积层共享时分解结果与参数对象会被复用避免重复计算与参数丢失通过_modify_data_type处理精度转换如float64分解结果转回原权重float16/float32时会先 clamp 到目标类型可表示范围防止溢出并通过_modify_memory_format保持与原始权重一致的contiguous_format/channels_last内存格式Torch 1.5 及以上版本生效。调用示例from amct_pytorch.tensor_decompose import auto_decomposition net Net() # 构建用户模型对象 net.load_state_dict(torch.load(src_path/weights.pth)) # 加载模型权重 net, changes auto_decomposition( # 执行张量分解 modelnet, decompose_info_pathdecomposed_path/decompose_info.json )示例要点先构建模型对象含预训练权重再执行分解changes中记录了分解前后层名映射可用于后续调试或权重迁移decompose_info.json即为decompose_network所需的分解信息文件建议以.json扩展名保存。使用注意事项[!NOTE] 说明当涉及模型训练时本接口的调用需在将模型参数传递给优化器之前如使用了torch.nn.parallel.DistributedDataParallelDDP则本接口的调用也需在将模型传递给 DDP 之前。本接口将原地修改传入的模型对象即分解后会改变用户传入的模型对象本身例外传入的模型是一个torch.nn.Conv2d对象该情况下本接口不会对其进行修改如发生分解则返回的模型是新构建的torch.nn.Module对象。从源码确认_replace_layers在is_conv2d时直接将新层容器赋值给返回的model新构建对象否则通过_relink_layers沿层名路径在原模型内逐级setattr/索引替换原地修改原模型结构。另外源码还提示了分布式训练场景的最佳实践tensor_decompose.py由于分解耗时较长更推荐先调用一次auto_decomposition获取分解信息文件并保存分解后的权重之后在每个训练进程中改用decompose_network依据文件重建结构并加载已保存的权重从而避免重复执行分解计算。与 decompose_network 的配合使用配套接口decompose_network接口文档的用法如下from amct_pytorch.tensor_decompose import decompose_network net Net() # 构建用户模型对象 net, changes decompose_network( # 加载分解信息文件将模型结构修改为张量分解后的结构 modelnet, decompose_info_pathdecomposed_path/decompose_info.json # 由auto_decomposition保存的分解信息文件路径 )使用decompose_network时需注意源码见 tensor_decompose.py用户输入的模型结构需与调用auto_decomposition获取分解信息文件时的模型结构一致分解信息文件必须与该模型结构配套使用decompose_network会通过_check_layer_info校验当前层与文件记录的in_channels、out_channels、kernel_size、stride、dilation、groups完全一致任一不匹配即抛出ValueError防止误用本接口仅对模型结构进行修改不会更新分解后的卷积权重权重的值为torch.nn.Conv2d()构建的默认值。如需 fine-tune请在调用auto_decomposition后将分解后的模型权重保存下来在调用decompose_network之后加载该权重再进行 fine-tune同样需要遵循优化器之前、DDP 之前调用与原地修改模型的约束。接口验证仓库测试用例 test_public_forwarding.py 中test_tensor_decompose_forwarding验证了公共导入路径amct_pytorch.tensor_decompose与真实实现路径amct_pytorch.classic.graph_based.amct_pytorch.tensor_decompose下auto_decomposition、decompose_network两个函数指向同一实现确保文档化导入方式可用def test_tensor_decompose_forwarding(self): tensor_decompose public path forwards to the real functions. base amct_pytorch.classic.graph_based.amct_pytorch.tensor_decompose for name in (auto_decomposition, decompose_network): self._assert_same(amct_pytorch.tensor_decompose, base, name, base)小结auto_decomposition为 AMCT 提供了开箱即用的 Conv2d 自动张量分解能力用户只需传入含预训练权重的torch.nn.Module模型即可得到分解后的模型、层名映射字典以及可复用的 JSON 分解信息文件。理解其分解模式FCFK/FCSK/SCFK/SCSK、设备与精度处理、共享权重处理等实现细节有助于在训练前正确编排分解 → 保存权重 → 多进程重建结构的完整流程同时避开优化器、DDP 封装等时序陷阱。若需将分解能力嵌入更完整的压缩流程可进一步参考 docs/zh/api/README.md 中的接口总览与 docs/zh/api/decompose_network.md 的配套文档。【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表