3j6<ddlmZddlZddlmZmZmZmZm Z m Z m Z ddl m Z mZmZmZmZmZmZmZmZddlmZddZddZddZ dd dZdS)!) annotationsN)AttributeProto FunctionProto GraphProto ModelProto NodeProtoSparseTensorProto TensorProto) make_attribute make_function make_graph make_model make_node make_tensormake_tensor_value_infoset_model_propstensor_dtype_to_np_dtype) from_arraynoder thresholdintvalue_constant_of_shapefloatreturnlist[NodeProto]c |jdkrtd|jd|jD].}|jdkrt d|jd|jdkr|j}|jd}|j}tj|tj }||kr|gcSttj t|tj | }t|j} tdg|g| } td |g|jttj |g|  } | | gcSt d |j|gS)zReplaces a Constant node with a large tensor (with more than threshold elements) by a sequence of nodes that produces a dummy constant of same shape as original tensor.Constantz!Node type must be 'Constant' not . sparse_valuezEThis feature is not yet implemented for a sparse constant (node name=).value__SHAPEdtypenamer!ConstantOfShapez'Replacement of constant with attribute )op_type TypeError attributer&NotImplementedErrortdimsnpprodint64rarraylistr data_typeroutput) rrrattr!new_namer.sizeinitr$ node_shapenew_nodes MC:\Projects\Historical\venv\Lib\site-packages\onnx/tools/replace_constants.py_replace_constantr=s |z!!MDLMMMNNN~   8~ % %%."i...  8w  EE*---H:D74rx000Dy  v bhtDzzBBBRRRD,U_==E"  J !!   +B*C5!Q!Q!QRR H ) ) ) )! Bch B B    6MonxGraphProto | FunctionProtoc t|ttfrt|j}n t dt |dt|D]0}t|jzt|j z1fd}tdg|dgd}tdg|dgd }i}t|D]x\}}|j d kr|jd}td |g||d g} td |j d| j d|j dg||dg} t|jd kr|jdjj} n t"j} td| j dg||dg| } td| j dg||dg| } td| j d| j dg||dg}td|j d|g|j }| | | | ||g||<zt'|dD]\}}||||d z<|d||d |t|tr.t-||j|j|j |j|jSt|tr.t5|j|j|j|j ||jSt dt |d)zReplaces all *ConstantOfShape* by node *Range* to avoid constant tensors. The function is not recursive. The recursivity is done by *replace_initializer_by_constant_of_shape*. Not implemented for type rc|vr||Sd} |d|}|vr||S|dz }()NT_)add)prefixir&existing_namess r< _find_namez9_replace_constant_of_shape_with_range.._find_nameYsr  ' '   v & & &M  ??q??D>))""4((( FA  r>rzeror) value_intonerFr( ReduceProd_NRange_RANGECast_RANGEf)to_NfDiv_FLATReshapeT)reverse initializersparse_initializer opset_imports) isinstancerrr3rr*typesetinputr5r enumerater)lenr+r-r4r FLOATsorteditemsinsertr r&r\r]r domain opset_import)r?nodesrrKcst0cst1updateinodeshapenarUaccldreshuprJs @r<%_replace_constant_of_shape_with_rangeryGs# M233BSX@DII@@@AAAUUN++#dj//)#dk***      Zjj&8&8%9Q G G GD Zjj&7&7%8A F F FD F ''00 t <, , ,  1  lUGjjE.F.F-G H H   [^QXa[$+a. 9 Z5((( ) ) *   t~  ! # #"$.BB"B v }zzU:K:K:K/L/L.MRT U U U v }zzU---/H/H.Ib Q Q Q  BIaL")A,/**___2M2M1N  QXa[%$8$+FFAr2q$/u FLLNND999&& r#%eeai   LLD LLD#z""   H I J"5     #}%%  J H I J *      <S <<< = ==r>ct|ttfrt|j}n t dt |di}t|D]\}}|jdkr|j dj }t|j |j dg|g}td|j|j}t#|j dj |} |j | |||<|D] \}} | ||< t|tr.t)||j |j|j|j|jSt|tr.t/|j|j |j|j||jSt dt |d) z7Replaces all fill value of all nodes *ConstantOfShape*.rBrr(rrFr'r[r^)r`rrr3rr*rardr)r+r-rr&r4rrcr5r appendrhr r\r]r rjrk) r?rrlrorprtensor new_tensorr;r6rxs r< _replace_constant_of_shape_valuer~s# M233BSX@DII@@@AAA F '' ! ! t <, , , "$ K)A31H0I  . DKHHT^A.3:FFF!!#&&& u \\^^ re #z""   H I J"5     #}%%  J H I J *      <S <<< = ==r>F?'FunctionProto | GraphProto | ModelProto ir_version int | None use_rangeboolc Z#ttrd}g}jD]^}|jdkrs tAdtCdd }g}tE#g}g}j#D]t}tI|j%}tMj'|tLj(}|kr||S|dz }|j d}|tStMj*tW|tLj(|tY|j-}t]d|g|j gtStMj*dg|}||#|j Bdkr<|t_|t`j1t |gvg}j2D]{}tI|j%}tMj'|tLj(}|kr||Rtgd|j4j d|j5j djD]}|jdkr?t |}t |dkr|dz }||Md}g}|j6D]}|j!tnj8krk|9drVt#|j:}tw|tw|j:krd}ty|j |}n|j!tnj=krng}|j>D]M}t#|} tw| tw|krd}|| N|rty|j |}|||rVt]|j|j |j }!|!j6|||!|dz }|||d kr\t|j #fd j D|zj ||!}"rt|"Sdkrt|"S|"SrtSdkrtSS)"aReplace initializers or constant node by nodes *ConstantOfShape* to reduce the size. This reduce the cost to write a unit test about a specific graph structure. Args: onx: ModelProto threshold: every initializer under this threshold is not impacted ir_version: initializer must be specified as input for `ir_version <= 3`, this must be specified if onx is :class:`FunctionProto` or :class:`GraphProto` use_range: if uses operator *Range* instead of *ConstantOfShape* to avoid constant tensors value_constant_of_shape: value to use as a value for all nodes *ConstantOfShape*, a high value may produce nan or inf predictions Returns: onx, modified ModelProto The function is designed so that the function can be reapplied on a modified model and either replace *ConstantOfShape* with *Range* operators, either replace the fill value for every *ConstantOfShape*. FrrDTr^rF)rrrrc Fg|]}t|pjS)rrrr)(replace_initializer_by_constant_of_shaper).0frr?rrrs r< z.sP     5#%7#(?        r>) functions producer_nameproducer_versionr doc_stringrj model_versionrc(i|]}|j|jS)keyr!)rps r< z.$sAAAaeQWAAAr>N z.Range was introduced in opset 11 but opset is r z7ConstantOfShape was introduced in opset 9 but opset is z-onx should be a GraphProto at this stage not r#r"r%r(rr'zKThis feature is not yet implemented for a sparse initializer (indices.name=z, values.name=r grc&g|] }|jv |Srr%)rrIremoveds r<rz.s% ; ; ;1QV7%:%:Q%:%:%:r>r[)@r`rrr)r=reextendr{r rjr&rcr5rkryr~rrgraphrrrrrrrmetadata_propsrrGversion RuntimeErrorrr*rarbr\tupler.r/r0r1rr2r3rr4rrr INT64r]r,indicesvaluesr+rGRAPHHasFieldridr GRAPHSgraphsr )$r?rrrrmodified new_nodesr cst_nodesnew_onx new_graph new_functionsmodelroimpop_setn_modificationsadditional_inputs new_initsr9r.r8r7r$new_sparse_initssp_init shape_nodesattsr6r new_graphssub_gnew_gr;rrs$````` @r<rrsS>#}%%%' H # #D|z))-dI?VWW y>>Q&&#H  +++   T " " " "  #   !. G F8== = "a ' '3C9PQQ Q #z"".< I!3S^$;             ]     #+ 1!3S^~:+     s! " "Q & &AAc.@AAAF E6 * * *  qqq !$ * *D'++--F{b  T\B%6%69%6"TT\TTT{b  T\A%5%5"<,0L<<<!KFM!\FNN c: & &VTS TTTUUUOIeeG#%ITYwt28,,, 9     T " " " 1i((( rxT "(;;;( K K K   )88  J YKRXse5999::      DI  !jAoo  $ $&x1BSYYKPP   13)   W\""wt28,,, 9    # #G , , , ! 5$_1 5 5">. 5 5 5   +#+# <: % %+D)=TUUK;1$$1$   [ ) ) ) >  Cx>///CLL4E4E/<E')',C a55BsuII%%#H(155C^222  Z - -ED"+#-"+0G E%yyBuII--#'%%e,,,,?(:>>C KK      # tz4;GGH   % %d + + +   X & & & q OO   T " " " "  H ; ; ; ; ; ; ;>O O J!/      @8?? ? "a ' '3E;RSS S :4S999!##/5LMMM Jr>)rrrrrrrr)r?r@rr@)r?r@rrrr@)rNFr) r?rrrrrrrrr) __future__rnumpyr/onnxrrrrrr r onnx.helperr r r rrrrrronnx.numpy_helperrr=ryr~rrr>r<rs|#"""""                      )(((((%%%%PO>O>O>O>d+>+>+>+>`!%( ^^^^^^^r>