报错 RuntimeError: Only consecutive 1-d tensor indices are supported in exporting aten::index_put to O
·
多个轴索引,存在多个数值,需要满足【:】所在轴的数值在内存中是连续的,也就是【:】只能出现在最后的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)
更多推荐


所有评论(0)