ONNX HardSwish

import sys
from pathlib import Path
ROOT = Path(".").resolve().parents[5]
# print(ROOT)
sys.path.extend([f"{ROOT}/tests"])
# from tools.tag_span import _create_span, _set_span, _verify_structural_equal_with_span
import tools
from d2py.utils.file import mkdir
root_dir = ".temp"
mkdir(root_dir )

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


class M(nn.Module):
    def forward(self, x):
        return x * (F.hardtanh(x + 3, 0., 6.) / 6.)

model = M()
model.eval()

shape = 1, 1000
xx = torch.rand(*shape, dtype=torch.float32, requires_grad=False)
# model = torch.jit.trace(model, xx)
# 导出模型
output_name = "test"
utils.export(
    model,               # torch 模型
    xx,                         # 模型输入或者对于多个输入，使用元组
    f"{root_dir}/{output_name}.onnx",               # 模型保存的位置（可以是文件或类似文件的对象）
    export_params=True,        # 将训练后的参数权重存储在模型文件内
    opset_version=17,          # 导出模型的 ONNX 版本
    do_constant_folding=True,  # 是否执行常量折叠以进行优化
    input_names = ['data'],    # 模型的输入名称
    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(%data : Float(1, 1000, strides=[1000, 1], requires_grad=0, device=cpu)):
  %/Constant_output_0 : Float(requires_grad=0, device=cpu) = onnx::Constant[value={3}, onnx_name="/Constant"](), scope: __main__.M:: # /tmp/ipykernel_444618/1751548276.py:9:0
  %/Add_output_0 : Float(1, 1000, strides=[1000, 1], requires_grad=0, device=cpu) = onnx::Add[onnx_name="/Add"](%data, %/Constant_output_0), scope: __main__.M:: # /tmp/ipykernel_444618/1751548276.py:9:0
  %/Constant_1_output_0 : Float(device=cpu) = onnx::Constant[value={0}, onnx_name="/Constant_1"](), scope: __main__.M:: # /media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/site-packages/torch/nn/functional.py:1551:0
  %/Constant_2_output_0 : Float(device=cpu) = onnx::Constant[value={6}, onnx_name="/Constant_2"](), scope: __main__.M:: # /media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/site-packages/torch/nn/functional.py:1551:0
  %/Clip_output_0 : Float(1, 1000, strides=[1000, 1], requires_grad=0, device=cpu) = onnx::Clip[onnx_name="/Clip"](%/Add_output_0, %/Constant_1_output_0, %/Constant_2_output_0), scope: __main__.M:: # /media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/site-packages/torch/nn/functional.py:1551:0
  %/Constant_3_output_0 : Float(requires_grad=0, device=cpu) = onnx::Constant[value={6}, onnx_name="/Constant_3"](), scope: __main__.M:: # /tmp/ipykernel_444618/1751548276.py:9:0
  %/Div_output_0 : Float(1, 1000, strides=[1000, 1], requires_grad=0, device=cpu) = onnx::Div[onnx_name="/Div"](%/Clip_output_0, %/Constant_3_output_0), scope: __main__.M:: # /tmp/ipykernel_444618/1751548276.py:9:0
  %output : Float(1, 1000, strides=[1000, 1], requires_grad=0, device=cpu) = onnx::Mul[onnx_name="/Mul"](%data, %/Div_output_0), scope: __main__.M:: # /tmp/ipykernel_444618/1751548276.py:9:0
  return (%output)