torch转ONNX：split

split() 将张量分割成块。每个块都是原始张量的视图。

如果 split_size_or_sections 是整数类型，那么张量将被等分成大小相等的块（如果可能的话）。如果张量沿着给定维度 dim 的大小不能被 split_size 整除，最后一个块将会更小。

如果 split_size_or_sections 是列表，那么张量将被分割成 len(split_size_or_sections) 个块，这些块在 dim 维度上的大小根据 split_size_or_sections 来确定。

import torch

torch.split?

Signature:
torch.split(
    tensor: torch.Tensor,
    split_size_or_sections: Union[int, List[int]],
    dim: int = 0,
) -> Tuple[torch.Tensor, ...]
Docstring:
Splits the tensor into chunks. Each chunk is a view of the original tensor.

If :attr:`split_size_or_sections` is an integer type, then :attr:`tensor` will
be split into equally sized chunks (if possible). Last chunk will be smaller if
the tensor size along the given dimension :attr:`dim` is not divisible by
:attr:`split_size`.

If :attr:`split_size_or_sections` is a list, then :attr:`tensor` will be split
into ``len(split_size_or_sections)`` chunks with sizes in :attr:`dim` according
to :attr:`split_size_or_sections`.

Args:
    tensor (Tensor): tensor to split.
    split_size_or_sections (int) or (list(int)): size of a single chunk or
        list of sizes for each chunk
    dim (int): dimension along which to split the tensor.

Example::

    >>> a = torch.arange(10).reshape(5, 2)
    >>> a
    tensor([[0, 1],
            [2, 3],
            [4, 5],
            [6, 7],
            [8, 9]])
    >>> torch.split(a, 2)
    (tensor([[0, 1],
             [2, 3]]),
     tensor([[4, 5],
             [6, 7]]),
     tensor([[8, 9]]))
    >>> torch.split(a, [1, 4])
    (tensor([[0, 1]]),
     tensor([[2, 3],
             [4, 5],
             [6, 7],
             [8, 9]]))
File:      /media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/site-packages/torch/functional.py
Type:      function

%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 __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 1, 1, 0, bias=False, groups=1)
    def forward(self, x):
        x = self.conv1(x)
        return torch.split(x, 8, dim=1) #, torch.split(x, [1, 3, 12], dim=1)

shape = 1, 3, 8, 8
x = torch.rand(*shape)

torch_model = Model()
# 导出模型
output_name = "split"
torch.onnx.export(
    torch_model,               # torch 模型
    x,                         # 模型输入或者对于多个输入，使用元组
    f"{temp_dir}/{output_name}.onnx",             # 模型保存的位置（可以是文件或类似文件的对象）
    export_params=True,        # 将训练后的参数权重存储在模型文件内
    opset_version=9,          # 导出模型的 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()

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()