o ;ήcF@sddlZddlZddlmZmZddlZddlZddlmZm Z ddlm Z ddl m Z eeZddZdd d Zdd d ZddZgdZGdddZ      dddZdddZdS)N)DictList)helper numpy_helper)onnx_pb)versioncCsdd|DS)z| Convert numpy float16 to python int. :param np_list: numpy float16 list :return int_list: python int list cSs.g|]}tt|dddddqS)HN)intbinviewzfill).0_rG/tmp/pip-target-vg8gfxp4/lib/python/onnxruntime/transformers/float16.py s.z%_npfloat16_to_int..r)np_listrrr_npfloat16_to_intsr"\o>@cCszdd}t|d||||}t|| |d| |}t|||td||}t|td|| | |}t|S)a? Convert float32 numpy array to float16 without changing sign or finiteness. Positive values less than min_positive_val are mapped to min_positive_val. Positive finite values greater than max_finite_val are mapped to max_finite_val. Similar for negative values. NaN, 0, inf, and -inf are unchanged. cSst||k||kSN)np logical_and)abcrrrbetween(sz&convert_np_to_float16..betweenrinfz-inf)rwherefloatfloat16)np_arraymin_positive_valmax_finite_valrrrrconvert_np_to_float16 s  r&cCst|tjstdt||jtjjkrOtjj|_|jr9t t |j||}t |}||j dd<g|jdd<|jrOt j|jdd}t |||}||_|S)aConvert tensor float to float16. Args: tensor (TensorProto): the tensor to convert. min_positive_val (float, optional): minimal positive value. Defaults to 1e-7. max_finite_val (float, optional): maximal finite value. Defaults to 1e4. Raises: ValueError: input type is not TensorProto. Returns: TensorProto: the converted tensor. 5Expected input type is an ONNX TensorProto but got %sNfloat32dtype) isinstance onnx_proto TensorProto ValueErrortype data_typeFLOATFLOAT16 float_datar&rarrayr int32_dataraw_data frombuffertobytes)tensorr$r% float16_dataint_list float32_list float16_listrrrconvert_tensor_float_to_float162s    r>cCst|j}t|j|j|Sr)rto_arrayshapermake_tensor_value_infonamer0)r9r@rrrmake_value_info_from_tensorWs rC)ArrayFeatureExtractor BinarizerCastMapCategoryMapperDictVectorizerFeatureVectorizerImputer LabelEncoderLinearClassifierLinearRegressor Normalizer OneHotEncoder SVMClassifier SVMRegressorScalerTreeEnsembleClassifierTreeEnsembleRegressorZipMapNonMaxSuppressionTopKRoiAlignResizeRangeCumSumMinMaxUpsamplec@s0eZdZdZdejfddZdejfddZdS) InitializerTrackerz'Class for keeping track of initializer. initializercCs||_g|_g|_dSr)r` fp32_nodes fp16_nodes)selfr`rrr__init__~s zInitializerTracker.__init__nodecCs$|r |j|dS|j|dSr)raappendrb)rcreis_node_blockedrrradd_nodeszInitializerTracker.add_nodeN) __name__ __module__ __qualname____doc__r,r-rd NodeProtorhrrrrr_{sr_Fc$s|dksJd|tttjjksJdd}|s2ttjtdkr2z ddl m } | }Wnwt |t j s@tdt||durFt}|durLg}t|}t|}td |d |d d |d |d|d|g} g} g} |dur|||}| |i} t}t}dd|jjD}dd|jjD}t trfdd|D}fdd|D}nsg}g}t|jjD]U\}}|j|vrdt|}|| |j<||jdt|}|jj}||||_t jj |jj!_"t#j$d|jg|gd|dg}|jj%&|| |||qt|jjD]V\}}|j|vrpdt|}|| |j<||jdt|}|jj}||||_t jj |jj!_"t#j$d|g|jgd|dg}|jj%&|| |||qi}| rg}| D]Z}t |t j r||jt |t j'rL|j(D]}|j)t jj*kr|j|vsJt+|||j<q|j%D]}|j|vrqt,t-|jD]}|j|| vr| |j||j|<qt,t-|jD]}|j|| vr| |j||j|<q|j.|vp|j|v}|jD]}||vr||/||q|r | |q|j.dkr>|j0D]}|jdkr<|j1dkr| j(| j?r |s t@dA| j=q| D]}!t,t-|!jD]V}|!j|}| D]K}"||"jkrk|jj}||"|!jd t|}||_t jj*|jj!_"|!jd!t|}t#j$d|g|gd|dg}|jj%&|||!j|<nq!qt,t-|!jD]V}|!j|}#| D]K}"|#|"jkr|jj}||"|!jd"t|}||_t jj*|jj!_"|!jd#t|}t#j$d|g|#gd|dg}|jj%&|||!j|<nqqvq|S)$aYConvert model tensor float type in the ONNX ModelProto input to tensor float16. Args: model (ModelProto): The ONNX model to convert. min_positive_val (float, optional): minimal positive value. Defaults to 5.96e-08. max_finite_val (float, optional): maximal finite value of float16. Defaults to 65504. keep_io_types (Union[bool, List[str]], optional): It could be boolean or a list of float32 input/output names. If True, model inputs/outputs should be left as float32. Defaults to False. disable_shape_infer (bool, optional): Skips running onnx shape/type inference. Useful if shape inference has been done. Defaults to False. op_block_list (List[str], optional): List of op types to leave as float32. Defaults to None, which will use `float16.DEFAULT_OP_BLOCK_LIST` as default. node_block_list (List[str], optional): List of node names to leave as float32. Defaults to None. force_fp16_initializers(bool): force converting all float initializers to float16. Default to false, which will convert only the one needed to avoid precision loss. Raises: ValueError: input type is not ModelProto. Returns: ModelProto: converted model. rzginvalid min_positive_val. smallest positive float16 value: subnormal 5.96e-08, and normalized 6.104e-05z4invalid max_finite_val. largest float16 value: 65504Nz1.2.0r) infer_shapesz4Expected model type is an ONNX ModelProto but got %sz"fp16 parameters: min_positive_val=z max_finite_val=z keep_io_types=z disable_shape_infer=z op_block_list=z node_block_list=z force_fp16_initializers=cS$g|]}|jjjtjjkr|jqSrr/ tensor_type elem_typer,r-r1rBrnrrrr$z,convert_float_to_float16..cSrorrprsrrrrrucg|]}|vr|qSrrrs keep_io_typesrrrcrvrrrsrwrrrrygraph_input_cast_graph_input_castCast )torBgraph_output_cast_graph_output_castr~ sequence_typezZinitializer is used by both fp32 and fp16 nodes. Consider add these nodes to block list:{} _input_cast_ _input_cast _output_cast_ _output_cast)Br!rfinfor"maxrparseonnx __version__onnx.shape_inferencernr+r, ModelProtor.r/DEFAULT_OP_BLOCK_LISTsetloggerdebugrfgraphinputoutputlist enumeraterBstradd value_infoCopyFromr-r2rqrrr make_nodereextend GraphProtor`r0r1r_rangelenop_typerh attributeiAttributeProtoggraphstr>tensors itertoolschainHasFieldritemsrbrCrainfoformat)$modelr$r%rxdisable_shape_infer op_block_listnode_block_listforce_fp16_initializersfunc_infer_shapernqueuevalue_info_list node_list name_mappinggraph_io_to_skipio_casts fp32_inputs fp32_outputsrrt output_name node_namenew_value_infonew_node input_namefp32_initializers next_levelqrgrattrkeyvaluererrrrwrconvert_float_to_float16s^  ,                                     C          rcCst|tjstdt||jtjjkrtdd}|jr$t |j}|j r/tj |j dd}|dur7t dt |||}tt|t|S)zSMeasure the maximum absolute difference after converting a float tensor to float16.r'z#Expected tensor data type is float.Nr(r)zexternal data not loaded!)r+r,r-r.r/r0r1r3rr4r6r7 RuntimeErrorr&amaxabsr()r9r$r% float32_datar:rrrfloat_to_float16_max_diffxs   r)rr)rrFFNNF)rloggingtypingrrnumpyrrrrrr, packagingr getLoggerrirrr&r>rCrr_rrrrrrs2    % o