torch 切片转换

%cd ../../..
import set_env
from d2py.utils.file import mkdir
temp_dir = ".temp"
mkdir(temp_dir)

/media/pc/data/lxw/ai/tvm-book/doc/tutorials/frontend

import torch
from torch import nn

class Model(nn.Module):
    def forward(self, x):
        return x[3:7]

shape = 192, 36
x = torch.rand(*shape)

torch_model = Model()
# 导出模型
output_name = "slice"
torch.onnx.export(
    torch_model,               # torch 模型
    x,                         # 模型输入或者对于多个输入，使用元组
    f"{temp_dir}/{output_name}.onnx",             # 模型保存的位置（可以是文件或类似文件的对象）
    export_params=True,        # 将训练后的参数权重存储在模型文件内
    opset_version=17,          # 导出模型的 ONNX 版本
    do_constant_folding=True,  # 是否执行常量折叠以进行优化
    input_names = ['data'],    # 模型的输入名称
    output_names = ['output'], # 模型的输出名称
    # dynamic_axes={'data' : {0 : 'batch_size'},    # 可变长度的轴
    #               'output' : {0 : 'batch_size'}}
)

import onnx
import tvm
from tvm import relay
onnx_model = onnx.load(f"{temp_dir}/{output_name}.onnx")
mod, params = relay.frontend.from_onnx(onnx_model, {"data": shape}, freeze_params=True)
with tvm.transform.PassContext(opt_level=3):
    mod = relay.quantize.prerequisite_optimize(mod, params)
mod.show()

def @main(%data: Tensor[(192, 36), float32] /* ty=Tensor[(192, 36), float32] span=/Slice.data:0:0 */) -> Tensor[(4, 36), float32] {
  strided_slice(%data, begin=[3i64], end=[7i64], strides=[1i64], axes=[0i64]) /* ty=Tensor[(4, 36), float32] span=/Slice:0:0 */
}

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(mod, params)
qmod.show()

def @main(%data: Tensor[(192, 36), float32] /* ty=Tensor[(192, 36), float32] span=/Slice.data:0:0 */) -> Tensor[(4, 36), float32] {
  strided_slice(%data, begin=[3i64], end=[7i64], strides=[1i64], axes=[0i64]) /* ty=Tensor[(4, 36), float32] span=/Slice:0:0 */
}