Skip to content

Commit

Permalink
postprocess_model() shifted above to fix the issue 23508 in opencv/op…
Browse files Browse the repository at this point in the history
…encv repository
  • Loading branch information
appsinfinity868 authored and asmorkalov committed Apr 27, 2023
1 parent 37e49e6 commit 6e7d3c4
Showing 1 changed file with 31 additions and 30 deletions.
61 changes: 31 additions & 30 deletions testdata/dnn/onnx/generate_onnx_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -524,6 +524,37 @@ def generate_slice_neg_starts():

generate_slice_neg_starts()

def postprocess_model(model_path, inputs_shapes):
onnx_model = onnx.load(model_path)

def update_inputs_dims(model, input_dims):
"""
This function updates the sizes of dimensions of the model's inputs to the values
provided in input_dims. if the dim value provided is negative, a unique dim_param
will be set for that dimension.
"""
def update_dim(tensor, dim, i, j, dim_param_prefix):
dim_proto = tensor.type.tensor_type.shape.dim[j]
if isinstance(dim, int):
if dim >= 0:
dim_proto.dim_value = dim
else:
dim_proto.dim_param = dim_param_prefix + str(i) + '_' + str(j)
elif isinstance(dim, str):
dim_proto.dim_param = dim
else:
raise ValueError('Only int or str is accepted as dimension value, incorrect type: {}'.format(type(dim)))

for i, input_dim_arr in enumerate(input_dims):
for j, dim in enumerate(input_dim_arr):
update_dim(model.graph.input[i], dim, i, j, 'in_')

onnx.checker.check_model(model)
return model

onnx_model = update_inputs_dims(onnx_model, inputs_shapes)
onnx.save(onnx_model, model_path)

input_2 = Variable(torch.randn(6, 6))
custom_slice_list = [
slice(1, 3, 1),
Expand Down Expand Up @@ -1916,36 +1947,6 @@ def forward(self, x):
model = GatherMultiOutput()
save_data_and_model("gather_multi_output", x, model)

def postprocess_model(model_path, inputs_shapes):
onnx_model = onnx.load(model_path)

def update_inputs_dims(model, input_dims):
"""
This function updates the sizes of dimensions of the model's inputs to the values
provided in input_dims. if the dim value provided is negative, a unique dim_param
will be set for that dimension.
"""
def update_dim(tensor, dim, i, j, dim_param_prefix):
dim_proto = tensor.type.tensor_type.shape.dim[j]
if isinstance(dim, int):
if dim >= 0:
dim_proto.dim_value = dim
else:
dim_proto.dim_param = dim_param_prefix + str(i) + '_' + str(j)
elif isinstance(dim, str):
dim_proto.dim_param = dim
else:
raise ValueError('Only int or str is accepted as dimension value, incorrect type: {}'.format(type(dim)))

for i, input_dim_arr in enumerate(input_dims):
for j, dim in enumerate(input_dim_arr):
update_dim(model.graph.input[i], dim, i, j, 'in_')

onnx.checker.check_model(model)
return model

onnx_model = update_inputs_dims(onnx_model, inputs_shapes)
onnx.save(onnx_model, model_path)

class UnsqueezeAndConv(nn.Module):
def __init__(self):
Expand Down

0 comments on commit 6e7d3c4

Please sign in to comment.