多个轴索引,存在多个数值,需要满足【:】所在轴的数值在内存中是连续的,也就是【:】只能出现在最后的dimension,不能出现在前面,先放到最后,然后用permute函数

错误的方式1:x[self.c1[:, 0], :, self.c1[:, 1], self.c1[:, 3], self.c1[:, 2]] = self.value

错误的方式1:x[self.c1[:, 0], self.c1[:, 1], self.c1[:, 3], :, self.c1[:, 2]] = self.value

正确的方式:

import onnx
import torch
import onnxsim
from torch import nn
from torch.nn import Module
from onnxsim import simplify

class nets(Module):
    def __init__(self):
        super(nets, self).__init__()
        self.c1 = torch.tensor([[0, 0, 1, 0], [1, 2, 2, 1]], dtype=torch.int64)
        self.value = torch.randn((2, 3), dtype=torch.float32)

    def forward(self, x):
        x[self.c1[:, 0], self.c1[:, 1], self.c1[:, 3], self.c1[:, 2], :] = self.value
        return x

if __name__=="__main__":
    path1 = r'/projects/Fisheye3D/tmp1.onnx'
    network = nets()
    x = torch.ones((3, 3, 3, 3, 3), dtype=torch.float32)
    cnt = 0
    for h in range(3):
        for i in range(3):
            for j in range(3):
                for k in range(3):
                    for w in range(3):
                        x[h, i, j, k, w] = cnt
                        cnt += 1
    network.eval()
    output = network(x)
    # [0, 0, 0, 0]
    # [0, 1, 0, 0]
    # [0, 2, 0, 0]
    # [2, 0, 2, 2]
    # [2, 1, 2, 2]
    # [2, 2, 2, 2]

    input_names = ["k1"]
    output_names = ["o1"]
    dynamic_axes = {
        '1': {0: 'batch_size', 1:"num", 2:'kk'},
        'o1': {0: 'number'}
    }
    with torch.no_grad():
        torch.onnx.export(
            network,
            x,
            path1,
            verbose = False,
            opset_version = 16,
            do_constant_folding = True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
            keep_initializers_as_inputs=True,
            # dynamic_axes = dynamic_axes,
            # input_names = input_names,
            # output_names = output_names,
        ) # grid_sampler argsort
    onnx_model = onnx.load(path1)
    # convert model
    model_simp, check = simplify(onnx_model)
    assert check, "Simplified ONNX model could not be validated"
    onnx.checker.check_model(onnx_model)
    onnx.save(model_simp, path1)
Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐