HardSigmoid

参考：HardSigmoid

HardSigmoid 函数

\[ \operatorname{HardSigmoid}(x) = \max(0, \min(1, \alpha x + \beta))， \]

通常情况 \(\alpha=0.2\) 和 \(\beta=0.5\)。

from pathlib import Path
temp_dir = Path(".temp")
temp_dir.mkdir(exist_ok=True)
model_path = f"{temp_dir}/HardSigmoid.onnx" # 模型存储路径

构建模型：

from onnxscript import opset20 as op
from onnxscript import ir
from onnxscript import script
from onnxscript import FLOAT
import onnx

@script()
def model(x: FLOAT[1, 3, 4, 4]) -> FLOAT[1, 3, 4, 4]:
    x = op.Add(x, x)
    # y = op.HardSigmoid(x)
    # x = op.Mul(x, y)
    # x = op.HardSigmoid(x)
    return x
    # return op.HardSwish(x,)

onnx.save_model(model.to_model_proto(), model_path,)

model: Already defined.

转换为 Relay 模型

from tvm.driver.tvmc.frontends import load_model

model = load_model(model_path)
model.mod.show()

def @main(%x: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=n0.x:0:0 */) -> Tensor[(1, 3, 4, 4), float32] {
  add(%x, %x) /* ty=Tensor[(1, 3, 4, 4), float32] span=n0:0:0 */
}

构建模板

from tvm.relay.dataflow_pattern import (
    TupleGetItemPattern, is_op, wildcard,
    is_constant, rewrite,
    DFPatternCallback
)
import tvm
from tvm.relay import op as _op
from tvm.relay import transform as _transform
from tvm import relay

def make_hard_sigmoid_pattern():
    r"""匹配 ONNX HardSigmoid 算子的模式"""
    x = wildcard()
    alpha = is_constant()
    x = is_op("multiply")(x, alpha)
    beta = is_constant()
    x = is_op("add")(x, beta)
    x = is_op("clip")(x)
    return x

compiler = "ccompiler"
pattern_table = [(f"{compiler}.hard_sigmoid", make_hard_sigmoid_pattern())]


seq = tvm.transform.Sequential(
    [
        _transform.MergeComposite(pattern_table),
        # _transform.AnnotateTarget(compiler),
        # _transform.MergeCompilerRegions(),
        # _transform.PartitionGraph(),
        # _transform.InferType(),
        # _transform.Inline(),
        # _transform.DefuseOps()
    ]
)
mod = seq(model.mod)

mod.show()

def @main(%x: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=n0.x:0:0 */) -> Tensor[(1, 3, 4, 4), float32] {
  add(%x, %x) /* ty=Tensor[(1, 3, 4, 4), float32] span=n0:0:0 */
}

import numpy as np

def get_calibration_dataset(mod, input_name):
    dataset = []
    input_shape = [int(x) for x in mod["main"].checked_type.arg_types[0].shape]
    for i in range(5):
        data = np.random.uniform(size=input_shape)
        dataset.append({input_name: data})
    return dataset

dataset = get_calibration_dataset(mod, "x")
with tvm.transform.PassContext(opt_level=3):
    with relay.quantize.qconfig(
        skip_conv_layers=[],
        # calibrate_mode="kl_divergence", 
        weight_scale="max",
        round_for_shift=True,
        # rounding="TONEAREST", # "UPWARD" or "TONEAREST"
        # calibrate_skip_layers=[],
        skip_dense_layer=False,
    ):
        qmod = relay.quantize.quantize(model.mod, model.params, dataset)
qmod.show()

def @main(%x: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=n0.x:0:0 */) -> Tensor[(1, 3, 4, 4), float32] {
  add(%x, %x) /* ty=Tensor[(1, 3, 4, 4), float32] span=n0:0:0 */
}