大象转身,TPU-MLIR适配DragGAN模型前向操作

draggan draggan是由谷歌、麻省理工学院和马克斯普朗克研究所创建的一种新的人工智能模型。
通过点击、拖动等简单的交互操作就能改变拍摄对象的姿势、形状和表情等。
draggan改变了传统的ps操作流程,只需简单拖拽起点和终点,ai就会根据图像的变化自动生成和补全图像。
draggan可处理的图像类型丰富多样,无论是人类表情的调整还是自然风景的变化,都可以在瞬息之内轻松实现。
draggan的全流程包含一个基于generator的前向操作和反向传播过程。本文主要介绍在tpu-mlir上适配draggan模型的前向操作的全部过程。
模型移植推理代码定位与模型导出适配的模型代码使用 xingangpan/draggan: official code for draggan (siggraph 2023) (github.com) ,模型的入口在 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 避开这一算子的使用。
对 conv/deconv filter 为动态输入情况的支持draggan 的模型结构中,有一部分 conv 和 deconv 的输入是固定权重,而 filterop 部份是动态的从上游计算得到的输入。这种情况在这之前未做考虑,需要添加支持。这包括在多个地方的代码更改。下面通过具体的报错提示来一步步分析、定位和解决。
model_transform 阶段在 tpu-mlir 的 converter 中,权重(weight)和 动态输入(dynamic input)存储在不同的变量中,其中,weight 通过 getweightop(name)  获取,input 通过 getoperand(name)  获取。如果不确定 op 是 dynamic input 还是 weight,可以使用 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 在适配模型的过程会经过多步转换和多次优化,为了保证转换后的正确性,tpu-mlir 会做三次正确性验证,分别针对 top dialect,tpu dialect 和 bmodel。top 和 tpu 层的正确性的核心代码位于 moduleinterpreter.[h/cpp],该过程会从输入开始,对每一个 op 分配空间,进行初始化(init),在初始化结束后进行推理(inference),并在最终对每个 op 进行析构(deinit)。而 deconv 的精度错误之一则来自于 inference 阶段时 init 和 inference 的分离。
在 init 时,deconv 会构造一个 dnnl 的实例,此时会直接 copy 一份 weight 在 dnnl 实例中,但由于该 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 添加这一 permute 的先决条件。确保该 permute 添加是针对 deconv 的动态权重,且同时不会重复添加。因此考虑在 deconv 的 operation 结构中添加 bool 类型的 dynweight_reorderd 参数。当 filter 不是 top.weightop (使用动态权重)且 dynweight_reordered 为 false (没有添加对动态 weight 的 permute)时,添加这一 permute,同时设置  dynweight_reorderd 参数为 true。
在 topops.td 文件对 deconv 添加 dynweight_reorderd  参数后,对 deconv 动态权重的图优化逻辑如下:
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 have 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 channle 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< mul 的结构,active 是 reducesum 操作:
因此基本可以确定是普通的 f16 溢出问题。验证 bf16,发现bf16 编译成功,进一步确认是溢出问题。
将这些层添加到 qtable 中,发现还是通过不了比对,值里面仍然会存在 inf。对比发现是在 active(reducesum) -> sqrt 的结构中间有两个 cast 导致的:
这两个 cast 没有作用,可以被优化掉,于是写图优化将这两个 cast 直接消除。优化后对应的 mlir 如下:
int8 也是相同的精度溢出问题,同样确认两个 cast 融合的操作能够覆盖 int8 的情况即可。
至此,draggan 适配的模型部份适配完成。
总结在一些情况下,在不影响结果的情况下直接修改模型的代码结构可以更容易的解决一些算子适配问题较大的模型测试错误定位到具体算子的情况下,优先考虑构建单侧尝试复现问题控制变量,设置对照,是缺少解决思路时寻找问题的一个较为通用的方案。 

浙江省科技厅厅长到访加速科技 Synology® 发布 DSM 7.1 Beta 测试版
人工智能为中国彩电行业插上了新翅膀
“差评”不断的iPhone12该不该换?
电压基准滤波器将32位ADC SNR提高6dB:ADI设计说明
浴霸触摸开关面板RISC-V MCU方案
大象转身,TPU-MLIR适配DragGAN模型前向操作
浅析行星齿轮箱在3D扫地机上的应用
瑞萨推出首款单芯片可扩展型优化智能手机平台MP5232
我国成全球电梯保有量最大的国家,急需发展电梯物联网
stratasys在大中华区公布打印材料降价和教育客户优惠活动
单片机控制电机原理以及程序怎么写
替换所有.v代码中字符串的方法
人工智能和自动化有着什么样的区别?
UWB技术怎么让Phone作为苹果汽车的钥匙
进入加密货币市场需要避免哪些隐患
5G赋能汽车 汽车行业正在经历惊人的创新
我国脑肿瘤注册登记研究平台上线,具有重大的学科价值和社会意义
锂电池自放电和过放电现象的研究
希捷为PS5和PS4游戏主机提供游戏硬盘_三星S22 Ultra开启移动影像新篇章
常见的hash算法有哪些及其原理是什么