经验分享使用 Rp 类对 pytorch 算子作替换操作

Posted 2021-09-15 极智视界

  本教程记录了使用 Rp 类对 pytorch 算子作替换操作。

  现在的 AI 训练框架众多，主流还是 tf 和 pytorch，部署框架也很多，如 TRT、Openvino、TVM、NCNN、Tengine。。。在训练和部署之间，往往还需要你再做一件事，那就是模型转换，特别是像海思、昇腾、比特大陆之类，都会优先支持 caffe，但是现在谁训练还在用 caffe 呢。

  ONNX 的出现就是为了解决以上的 Gap，但 ONNX 对于一些自研网络又往往做不到全覆盖。所以学习一些模型转换的技巧也是需要的。

  这里介绍一下 from_pytorch 进行模型转换时进行算子替换的方法，主要使用了 Rp 类来实现。

  Rp 类的定义如下：

class Rp(object):
    def __init__(self, raw, replace, **kwargs):
        self.obj = replace
        self.raw = raw

    def __call__(self, *args, **kwargs):
        if not NET_INITTED:
            return self.raw(*args, **kwargs)
        for stack in traceback.walk_stack(None):
            if 'self' in stack[0].f_locals:
                layer = stack[0].f_locals['self']
                if layer in layer_names:
                    log.pytorch_layer_name = layer_names[layer]
                    print('984', layer_names[layer])
                    break
        out = self.obj(self.raw, *args, **kwargs)
        return out

  以 sigmoid 算子为例，使用 Rp 类对 torch.sigmoid 进行替换操作的姿势如下：

torch.sigmoid = Rp(torch.sigmoid, _sigmoid)

  其中 _sigmoid 你可以自己根据目标框架的需求格式进行自定义，_sigmoid 替换了原生 torch.sigmoid 的实现。

  例如，我这里的目标是将 pytorch 模型转换为结构清晰的 json 格式的模型结构文件，以备推理框架的调用，可以像下面这么来写_sigmoid：

def _sigmoid(raw, input):
    x = raw(input)
    INLINE = True
    name = log.add_layer(name="sigmoid_")
    log.add_blobs([x], name=name)
    LayerOut_id[int(id(x))] = name

    # no weight to extract
    # add json params
    sigmoid_params = dict({"layerStyle": "activation",
                           "layerName": name,
                           "inputName": log.blobs(input),
                           "active_type": "kSIGMOID"})

    js.data['network'].append(sigmoid_params)
    INLINE = False
    return x

  收工~

扫描下方二维码即可关注我的微信公众号【极智视界】，获取更多AI经验分享，让我们用极致+极客的心态来迎接AI ！