torch.index_put_函数

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

class net(Module):
    def __init__(self):
        super().__init__()
        self.c1 = torch.tensor([0, 2], dtype=torch.int32, device='cuda')
        self.c2 = torch.tensor([0, 2], dtype=torch.int32, device='cuda')
        self.c3 = torch.tensor([0, 2], dtype=torch.int32, device='cuda')

    def forward(self, x):
        r = x[self.c1, :, self.c2, self.c3]
        return r

if __name__=="__main__":
    path1 = r'/projects/Fisheye3D/tmp1.onnx'
    network = net().cuda()
    x = torch.ones((3, 3, 3, 3)).cuda()
    cnt = 0
    for i in range(3):
        for j in range(3):
            for k in range(3):
                for w in range(3):
                    x[i, j, k, w] = cnt
                    cnt += 1
    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'}
    }
    torch.onnx.export(
        network,
        x,
        path1,
        verbose = False,
        opset_version = 17,
        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)