Ngg ddlmZmZmZmZddlZddlmZmZdedee eefdee eefdefdZ dS) )AnyDictListSetN) ModelProtoValueInfoProtomodel input_dims output_dimsreturnc  t dttdttddfd}| |jj| |jj| |jjdtdtdtd tddf fd }|jjD]4}|j }||}t|D]\}} ||| ||5|jjD]4} | j } || } t| D]\}} || | || 5tj||S) aMThis function updates the dimension sizes of the model's inputs and outputs to the values provided in input_dims and output_dims. if the dim value provided is negative, a unique dim_param will be set for that dimension. Example. if we have the following shape for inputs and outputs: * shape(input_1) = ('b', 3, 'w', 'h') * shape(input_2) = ('b', 4) * shape(output) = ('b', 'd', 5) The parameters can be provided as: :: input_dims = { "input_1": ['b', 3, 'w', 'h'], "input_2": ['b', 4], } output_dims = { "output": ['b', -1, 5] } Putting it together: :: model = onnx.load('model.onnx') updated_model = update_inputs_outputs_dims(model, input_dims, output_dims) onnx.save(updated_model, 'model.onnx') dim_param_set value_infosr Nc|D]L}|jjj}|jD]1}|dr||j2MdS)N dim_param)type tensor_typeshapedimHasFieldaddr)rrinforrs X/var/www/html/ai-engine/env/lib/python3.11/site-packages/onnx/tools/update_model_dims.pyinit_dim_param_setz6update_inputs_outputs_dims..init_dim_param_set0sk  5 5DI)/Ey 5 5<< ,,5!%%cm444 5 5 5tensorrjnamec |jjjj|}t |t r|dkrJ|dr,|j|kr!td|d|d|d|jd ||_dS|dzt|z}|vrtd |d|d ||_ dSt |tr ||_ dStd t|) Nr dim_valuez!Unable to set dimension value to z for axis z of z'. Contradicts existing dimension value ._z-Unable to generate unique dim_param for axis z,. Please manually provide a dim_param value.z@Only int or str is accepted as dimension value, incorrect type: ) rrrr isinstanceintrr ValueErrorstrr)rrrr dim_protogenerated_dim_paramrs r update_dimz.update_inputs_outputs_dims..update_dim=szK+15a8 c3   axx%%k22y7Jc7Q7Q$VCVV1VVRVVV@I@SVVV'* ###&*Sj3q66&9#&-77$BBBtBBB': ### S ! ! "%I   ^SWX[S\S\^^ r)setrr&rrgraphinputoutput value_inforr$r enumerateonnxchecker check_model)r r r rr)input_ input_name input_dim_arrrrr- output_nameoutput_dim_arrrs @rupdate_inputs_outputs_dimsr8 sF"eeM53x5.2>.B5 5555}ek&7888}ek&8999}ek&<===>34.+#33[ ":.  .. 3 3FAs JvsAz 2 2 2 2 3+$44k $[1// 4 4FAs JvsA{ 3 3 3 3 4 LU### Lr) typingrrrr onnx.checkerr0rrr&r8rrr<s ('''''''''''++++++++V VS$s)^$Vc49n%V VVVVVVr