HardTanh Relay 实现

%cd ..
import numpy as np
import set_env
from d2py.utils.file import mkdir
root_dir = ".temp"
mkdir(f"{root_dir}/logs")

/media/pc/data/lxw/ai/tvm-book/doc/dev/ops
ROOT: /media/pc/data/lxw/ai/tvm-book

import torch
from torch.nn import functional as F
from torch import nn
from torch.onnx import OperatorExportTypes, utils

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 1, 1, bias=False)
        self.hard_tanh = nn.Hardtanh(-2, 2)

    def forward(self, x):
        x1 = self.hard_tanh(x)
        x2 = (F.hardtanh(x + 3, 0., 6.) / 6.) # 等价于 hard_sigmoid(x)
        return x1, x2

model = M()
model.eval()

shape = 1, 3, 8, 8
input_name = "x"
xx = torch.rand(*shape, dtype=torch.float32, requires_grad=False)
# model = torch.jit.trace(model, xx)
# 导出模型
output_name = "hard-tanh"
utils.export(
    model,               # torch 模型
    xx,                         # 模型输入或者对于多个输入，使用元组
    f"{root_dir}/{output_name}.onnx",               # 模型保存的位置（可以是文件或类似文件的对象）
    export_params=True,        # 将训练后的参数权重存储在模型文件内
    opset_version=9,          # 导出模型的 ONNX 版本
    do_constant_folding=True,  # 是否执行常量折叠以进行优化
    input_names = [input_name],    # 模型的输入名称
    output_names = ['output'], # 模型的输出名称
    keep_initializers_as_inputs=True,
    # export_modules_as_functions=True,
    verbose=True,
    operator_export_type=OperatorExportTypes.ONNX_FALLTHROUGH,
    # dynamic_axes={'data' : {0 : 'batch_size'},    # 可变长度的轴
    #               'output' : {0 : 'batch_size'}}
)

Exported graph: graph(%x : Float(1, 3, 8, 8, strides=[192, 64, 8, 1], requires_grad=0, device=cpu)):
  %output : Float(1, 3, 8, 8, strides=[192, 64, 8, 1], requires_grad=0, device=cpu) = onnx::Clip[max=2., min=-2., onnx_name="/hard_tanh/Clip"](%x), scope: __main__.M::/torch.nn.modules.activation.Hardtanh::hard_tanh # /media/pc/data/lxw/envs/anaconda3x/envs/py312/lib/python3.12/site-packages/torch/nn/functional.py:1551:0
  %/Constant_output_0 : Float(requires_grad=0, device=cpu) = onnx::Constant[value={3}, onnx_name="/Constant"](), scope: __main__.M:: # /tmp/ipykernel_2617739/3162195591.py:14:0
  %/Add_output_0 : Float(1, 3, 8, 8, strides=[192, 64, 8, 1], requires_grad=0, device=cpu) = onnx::Add[onnx_name="/Add"](%x, %/Constant_output_0), scope: __main__.M:: # /tmp/ipykernel_2617739/3162195591.py:14:0
  %/Clip_output_0 : Float(1, 3, 8, 8, strides=[192, 64, 8, 1], requires_grad=0, device=cpu) = onnx::Clip[max=6., min=0., onnx_name="/Clip"](%/Add_output_0), scope: __main__.M:: # /media/pc/data/lxw/envs/anaconda3x/envs/py312/lib/python3.12/site-packages/torch/nn/functional.py:1551:0
  %/Constant_1_output_0 : Float(requires_grad=0, device=cpu) = onnx::Constant[value={6}, onnx_name="/Constant_1"](), scope: __main__.M:: # /tmp/ipykernel_2617739/3162195591.py:14:0
  %7 : Float(1, 3, 8, 8, strides=[192, 64, 8, 1], requires_grad=0, device=cpu) = onnx::Div[onnx_name="/Div"](%/Clip_output_0, %/Constant_1_output_0), scope: __main__.M:: # /tmp/ipykernel_2617739/3162195591.py:14:0
  return (%output, %7)

备注

HardTanh(x, min_val, max_val) 函数在 ONNX 和 Relay 中均使用 clip(x, min_val, max_val) 替代。

import tvm
from tvm import relay
data_np = (np.random.randint(0, 256, shape)/255).astype("float32")
data_torch = torch.from_dlpack(data_np)

model = M().eval()
scripted_model = torch.jit.trace(model, data_torch).eval()
shape_list = [(input_name, shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
tvm.IRModule.from_expr(mod["main"]).show()

def @main(%x: Tensor[(1, 3, 8, 8), float32] /* span=aten::hardtanh_0.x:0:0 */) {
  %0 = add(%x, 3f /* span=aten::add_0:0:0 */) /* span=aten::add_0:0:0 */;
  %1 = clip(%0, a_min=0f, a_max=6f) /* span=aten::hardtanh_1:0:0 */;
  %2 = clip(%x, a_min=-2f, a_max=2f) /* span=aten::hardtanh_0:0:0 */;
  %3 = divide(%1, 6f /* span=aten::div_0:0:0 */) /* span=aten::div_0:0:0 */;
  (%2, %3)
}

import tvm
from tvm import relay
import onnx
onnx_model = onnx.load(f"{root_dir}/{output_name}.onnx")
mod, params = relay.frontend.onnx.from_onnx(onnx_model, {input_name: shape})
tvm.IRModule.from_expr(mod["main"]).show()

def @main(%x: Tensor[(1, 3, 8, 8), float32] /* ty=Tensor[(1, 3, 8, 8), float32] span=/hard_tanh/Clip.x:0:0 */) -> (Tensor[(1, 3, 8, 8), float32], Tensor[(1, 3, 8, 8), float32]) {
  %0 = add(%x, 3f /* ty=float32 span=/Constant:0:0 */) /* ty=Tensor[(1, 3, 8, 8), float32] span=/Add:0:0 */;
  %1 = clip(%0, a_min=0f, a_max=6f) /* ty=Tensor[(1, 3, 8, 8), float32] span=/Clip:0:0 */;
  %2 = clip(%x, a_min=-2f, a_max=2f) /* ty=Tensor[(1, 3, 8, 8), float32] span=/hard_tanh/Clip:0:0 */;
  %3 = divide(%1, 6f /* ty=float32 span=/Constant_1:0:0 */) /* ty=Tensor[(1, 3, 8, 8), float32] span=/Div:0:0 */;
  (%2, %3) /* ty=(Tensor[(1, 3, 8, 8), float32], Tensor[(1, 3, 8, 8), float32]) */
}