TPU-MLIR适配DragGAN模型前向操作实现_AI热点日报

DragGAN适配TPU-MLIR时，将RandomNormalLike算子设noise_mode=const规避不支持问题；对动态权重Conv DeConv调整filter获取逻辑、添加Permute重排，并将Dnnl实例设置移至inference阶段以解决精度错误。

DragGAN 模型适配 TPU-MLIR 教程

本教程将详细介绍如何将 DragGAN 模型适配到 TPU-MLIR 平台上。DragGAN 是由谷歌、麻省理工学院和马克斯普朗克研究所创建的一种新型人工智能模型，通过简单的点击和拖动操作即可改变对象的姿势、形状和表情，大幅简化传统 PS 流程。我们将重点讲解 DragGAN 模型在 TPU-MLIR 上的前向操作适配过程，包括模型导出、算子适配以及常见错误的解决思路。

提示： 开始前请确保已熟悉 TPU-MLIR 的基本使用方法，并拥有可运行的 DragGAN 源码环境。

模型移植

推理代码定位与模型导出

适配的模型代码使用 XingangPan/DragGAN (SIGGRAPH 2023)。模型的入口位于 DragGAN/viz/renderer.py:357。在这里可以直接引入 TPU-MLIR 提供的 gen_shell 工具，通过 trace 生成 workspace 文件夹、onnx/pt 模型以及默认的转换脚本：

from utils.gen_shell import generate

generate(
    "DragGan",
    G,
    dict(
        ws=ws,
        c=label
    ),
    "../draggan_workspace",
)

运行源码 README.md 中提供的脚本 python visualizer_drag_gradio.py，成功执行后可在同级目录得到如下目录结构：

draggan_workspace
├── cali_data
│   └── data.npz
├── convert.sh
├── DragGan.onnx
├── DragGan.pt
├── data.npz
└── cali_data

模型移植过程中错误的分析和解决

RuntimeError: Op not support:{'RandomNormalLike'}

在 model_transform 阶段，发现存在不支持的算子 RandomNormalLike：

原因分析： RandomNormalLike（随机数相关）算子 1684x 无法支持，必须尝试在原模型中避开这些算子。定位到模型代码，发现该算子用于提供一个噪音供下游使用。源码提供了三种噪音生成方式：random（随机噪音）、const（常量噪音）和 none（不提供噪音）。通过设置 noise_mode = const 可避开该算子。

提示： 遇到不支持的算子时，首先查看模型源码是否提供了可选参数（如 noise_mode），这往往是快速解决问题的路径。

对 Conv/DeConv Filter 为动态输入情况的支持

DragGan 模型结构中，部分 Conv 和 DeConv 的输入是固定权重，而 FilterOp 部分是动态从上游计算得到的输入。这种情况先前未做考虑，需要添加支持。下面通过报错提示逐步分析、定位和解决。

model_transform 阶段

在 tpu-mlir 的 Converter 中，权重（weight）和动态输入（dynamic input）存储在不同的变量中：

weight 通过 getWeightOp(name) 获取
input 通过 getOperand(name) 获取
不确定时 可使用 getOp(name) 获取

对 DragGan 运行 model_transform.py 时，会遇如下报错：KeyError: '/synthesis/b8/conv0/Transpose_output_0'

对应模型结构，发现该 DeConv 的输入 /synthesis/b8/conv0/Transpose_output_0 是作为 Weight 获取的。

解决方法：将 ConvTranspose 的 filter_opd 获取逻辑改为 getOp。

同理，另一个 KeyError 中 DeConv 的 filter 来自动态输入，将 DeConv 获取 filter 结点的逻辑也改为 getOp。

在 model_transform 阶段，模型首先转换到 DragGAN_origin.mlir，再经过 --shape-infer、--canonicalize 等过程，转换为可通过 model_runner.py 推理的 Top Dialect 描述的 mlir 文件。对 Top 层做推理验证时，DragGan 模型报出精度为零的错误。观察错误信息，发现是在 DeConv 层之后精度出现问题，且仅在 DeConv 的 filter 是动态输入时发生。

构建了 filter 是动态输入的 DeConv 作为单侧，成功复现该错误：

class DeConvCase(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.deconv = nn.ConvTranspose2d(4, 4, [2, 2], stride=[1, 1], bias=False)
        self.deconv.weight.data = weight
    def forward(self, x, y):
        output_padding = self.deconv._output_padding(
            x,
            None,
            [2, 2],
            [0, 0],
            [2, 2], 
            1,
            [1, 1],
        ) 
        out = F.conv_transpose2d(x, y, None, [1, 1], 0, output_padding, 1, 1)
        return out, self.deconv(x)

通过断点调试，发现错误原因有两个：

正确性验证阶段推理时，在 init() 时设置了权重，但此时 weight 还没有设置；
动态输入时没有做对应的权重重排（WeightReorder）。

tpu-mlir 在适配过程中会经过多步转换和多次优化，为保证正确性，会做三次正确性验证（Top Dialect、Tpu Dialect、bmodel）。Top 和 Tpu 层的正确性核心代码位于 ModuleInterpreter.[h/cpp]，该过程从输入开始，对每个 Op 分配空间进行初始化（init），结束后推理（inference），最后析构（deinit）。DeConv 的精度错误之一来自 Inference 阶段时 init 和 inference 的分离。

在 init 时，DeConv 构造 Dnnl 实例，此时会直接 copy 一份 Weight。由于 filter 为动态输入，init 时值尚未传入，传入的 filter 值为全零，导致 inference 阶段出现错误。解决方法：将 init 过程中对 Dnnl 实例的 setup 移到 inference 阶段即可。Conv 也有同样问题，修改逻辑相同。

对 onnx 模型，DeConv 的 filter 权重存储方式是 input channel first（shape 为 [ic, oc, kw, kh]），而后端计算大多需要 output channel first（[oc, ic, kw, kh]）。在 OnnxConverter 中，原本对 DeConv 的权重会进行一次转置：

而动态权重无法实现该操作。因此，需要添加一个图优化：当 DeConv 的 filter 是动态时，在其前面添加一个 [oc, ic] 互换的 Permute 操作。添加 Permute 时需注意先决条件，确保只针对 DeConv 的动态权重且不会重复添加。在 DeConv 的 Operation 结构中添加 bool 类型参数 dynweight_reorderd。当 filter 不是 top.WeightOp （使用动态权重）且 dynweight_reordered 为 false 时，添加 Permute，同时设置该参数为 true。

在 TopOps.td 文件对 DeConv 添加 dynweight_reorderd 参数后，图优化逻辑如下：

struct ReorderDynWeight : public OpRewritePattern {
  using OpRewritePattern::OpRewritePattern;
  LogicalResult matchAndRewrite(DeconvOp op,
                                PatternRewriter &rewriter) const override {
    auto filter_shape = module::getShape(op.getFilter()); //  or 
    if (module::isWeight(op.getOperand(1))) {
      return failure();
    }
    bool dyn_weight_reorderd = op.getDynweightReorderd();
    if(dyn_weight_reorderd){
      return failure();
    }
    if (isa(op.getOperand(1).getDefiningOp())) {
      auto permute_op =
          dyn_cast(op.getOperand(1).getDefiningOp());
      // erase if already ha ve this permute but from original graph
      std::vector ps = {1, 0, 2, 3};
      auto order = module::getI64Array(permute_op.getOrder());
      if (*order == ps) {
        permute_op.replaceAllUsesWith(permute_op.getInput());
        rewriter.eraseOp(permute_op);
        op.setDynweightReorderd(true);
        return success();
      }
    }
    rewriter.setInsertionPointAfterValue(op.getFilter());
    std::string name = module::getName(op.getOutput()).str();
    auto loc =
        NameLoc::get(rewriter.getStringAttr(name + "_reorder_permute"));
    std::vector order = {1, 0};
    auto filter_dim = filter_shape.size();
    for (int i = 2; i < filter_dim; i++) {
      order.push_back(i);
    }
    auto p_type =
        UnrankedTensorType::get(module::getElementType(op.getFilter()));
    std::vector attrs;
    attrs.emplace_back(
        rewriter.getNamedAttr("order", rewriter.getI64ArrayAttr(order)));
    auto new_permute_op = rewriter.create(
        loc, p_type, ValueRange{op.getFilter()}, attrs);
    new_permute_op.shape_inference();
    op.setOperand(1, new_permute_op.getOutput());
    op.setDynweightReorderd(true);
    return success();
  }
};

这里做了一个额外判断：当 DeConv 的 filter 位置已经是 Permute 且其 order 与待添加的 Permute 相同（1,0,2,3）时，两个 Permute 可直接融合，删除该 Permute 并返回。其他情况则插入一个额外的 Permute 操作。Conv 层同样要支持动态 weight 的权重重排，需添加相同的图优化。

另外，Top 层的 shape-infer 早于图优化，因此做 shape-infer 时动态 weight 的 shape 仍是 input channel first，所以 DeConv 的 output_shape 的 dim[1] 应基于 filter_shape[1] 判断。对应修改位于 lib/Dialect/Top/Interfaces/Deconv.cpp：

bmodel 运行错误

ASSERT /workspace/nntoolchain/TPU1686/bm1684x/cmodel/src/cmodel_common.cpp: gather_data: 207: dst_offset < (1<<18)

在大模型中定位此错误较难，可通过 mlir_cut.py 逐步缩小范围，得到最小可复现的 mlir：

mlir_cut.py --mlir *tpu.mlir --output_names /synthesis/b64/conv0/Conv_output_0_Conv --input_names /synthesis/b32/conv1/Mul_3_output_0_Mul,/synthesis/b64/conv0/Reshape_3_output_0_Reshape
tpuc-opt DragGan_bm1684x_f32_final.mlir --codegen="model_file=DragGan_f32.bmodel embed_debug_info=true" -o /dev/null
model_runner.py --input fake_data.npz --model DragGan_f32.bmodel --output DragGan_bm1684x_f32_model_outputs.npz

进一步构建了可复现该错误的单元测试：

通过控制变量，得到以下现象：

关闭 layer-group，模型运行正常不报错 → 问题基本出在 tpu-mlir 部分而非后端算子。
将 DeConv 的 filter 从动态改为静态，模型运行正常 → 问题仍是动态 Weight 导致。
构建基本的 DeConv 算子（静态和动态）均运行正常，与上述单侧对比发现区别在于单个 DeConv 算子不会进行 LayerGroup → 问题定位到 tpu-mlir 的 LayerGroup 部分代码。

进一步对比正常和出错的 final.mlir，发现 dynamic weight 和 weight 的 slice 属性不一致：

top.Weight 的 layer-group 比较特殊：top.Weight 在整个 layer-group 都保存在 local memory 中（hold_in_lmem = true）；同时 weight 也不能切分 slice，每个 slice 都要用到完整的 filter，导致结果错误。

解决方法：单独针对 dynamic weight 处理，包括设置其生命周期（hold_in_mem = true），以及将其 slice 设置为长度为 1、元素为其 shape 对应维度值的列表。该过程可在 lib/Dialect/Tpu/Transforms/LayerGroup/LayerGroupUtil.cpp 的 backward_update_slice 方法中完成：

优化后再次对比两个单例的 final.mlir，发现 dynamic weight 的 slice 信息已与普通 weight 完全相同：

F16 和 int8 精度问题

在解决 F32 的 bug 后，F16 和 int8 的 tpu 层 mlir 仍存在精度问题。原以为是 DeConv 的 F16 适配问题，通过使用 mlir_debugger 对每一层用正确数值推理（或观察输出的 npz 文件以及 npz_tool 比对结果），发现出错的是 Active → Mul 结构，其中 Active 是 ReduceSum 操作：

因此基本确定是普通的 F16 溢出问题。验证 BF16，发现 BF16 编译成功，进一步确认是溢出问题。

将这些层添加到 qtable 中，发现仍无法通过比对，值中仍然存在 inf。对比发现是在 Active(ReduceSum) → Sqrt 的结构中间有两个 Cast 导致：

这两个 Cast 没有作用，可被优化掉。于是编写图优化将这两个 Cast 直接消除。优化后对应的 mlir 如下：

int8 也是相同精度溢出问题，确认两个 Cast 融合的操作能覆盖 int8 的情况即可。

提示： 遇到精度问题时，使用 mlir_debugger 逐层比对输出，可快速定位异常层。特别注意 F16/BF16 的取舍和 Cast 累加带来的溢出风险。

常见问题

问题： 模型导出时生成的工作空间缺少某些文件？
答案： 确保已完整克隆 DragGAN 源码并正确安装依赖。运行脚本前，检查 visualizer_drag_gradio.py 中是否已导入 gen_shell 工具，并确认路径参数正确。
问题： 转换过程中间出现 KeyError，提示找不到某个层的输出？
答案： 通常是因为该层的 filter 是动态输入，但 Converter 仍使用 getWeightOp 查找。按照本文所述，将该处改为 getOp 即可。
问题： bmodel 运行时出现 gather_data 断言错误？
答案： 该错误多半由动态 Weight 的 slice 设置不当引起。请检查 LayerGroup 部分是否已按本文方法修改 backward_update_slice，确保动态 weight 的 slice 与普通 weight 一致。
问题： F16/int8 精度始终无法达标，且出现 inf 或 NaN？
答案： 可能是中间 Cast 重复导致数值溢出。可尝试添加图优化消除冗余 Cast 层，或将部分层切换到 BF16 避免溢出。

总结

在不影响结果的前提下，直接修改模型代码结构（如设置 noise_mode = const）可以更简单解决算子适配问题。
对于较大模型中的错误，优先构建最小单侧（如单层 DeConv）尝试复现问题，能大幅缩小排查范围。
控制变量、设置对照（如关闭 layer-group、切换静态/动态权重）是缺乏解决思路时寻找问题的通用方案。

本教程完整覆盖了 DragGAN 在 TPU-MLIR 上的适配流程，从模型导出到常见错误修复，希望能帮助开发者快速上手。若您在适配过程中遇到其他问题，欢迎在 TPU-MLIR 社区交流讨论。