U T?h@sddlZddlZddlZddlZddlZddlZddlZddlmZddl m Z m Z m Z m Z ddlZddlZddlZddlmZddlmZddlmZddlmZddlmZdd lmZdd lmZdd lm Z m!Z!m"Z"m#Z#e$e%Z&d d dddgZ'ej(dej)dej*diZ+Gddde"Z,Gddde#Z-Gddde"Z.Gddde!Z/Gddde!Z0e/ddfe0dd fe.d!dfd"Z1Gd#d$d$Z2Gd%d&d&Z3dS)'N)Path)DictListTupleUnion) Precision)float_to_float16_max_diff) FusionOptions)IOBindingHelper) OnnxModel)optimize_model)torch_onnx_export) GPT2ConfigGPT2LMHeadModel GPT2Model TFGPT2ModelZ distilgpt2gpt2z gpt2-mediumz gpt2-largezgpt2-xlMb@?g?g@cs,eZdZdZfddZfddZZS)GPT2ModelNoPastState2Here we wrap a class to disable past state output.cst|dSNsuper__init__selfconfig __class__b/var/www/html/venv/lib/python3.8/site-packages/onnxruntime/transformers/models/gpt2/gpt2_helper.pyr+szGPT2ModelNoPastState.__init__cstj|dddS)NF) use_cache return_dict)rforwardr input_idsrrr r#.szGPT2ModelNoPastState.forward__name__ __module__ __qualname____doc__rr# __classcell__rrrr r(s rcs,eZdZdZfddZfddZZS)TFGPT2ModelNoPastStatercsd|_t|dS)NF)r!rrrrrr r5szTFGPT2ModelNoPastState.__init__cstj|ddS)NF)r!)rcallr$rrr r#9szTFGPT2ModelNoPastState.forwardr&rrrr r,2s r,cs8eZdZdZfddZeddZfddZZS) MyGPT2ModelzMHere we wrap a class for Onnx model conversion for GPT2Model with past state.cst|dSrrrrrr r@szMyGPT2Model.__init__cCst|ddttfrt|d|kr:t|dddks>tg}t|D]@}|tj|d|d d|d|d dfddqJ|dt|fS|S)Nr)dim) isinstancetuplelistlenAssertionErrorrangeappendtorchcatZ unsqueeze)result num_layerZpresentirrr post_processCs( *zMyGPT2Model.post_processcs&tj||||dd}t||jjSNF) position_idsattention_maskpast_key_valuesr"rr#r.r>rn_layerrr%r@rApastr;rrr r#VszMyGPT2Model.forward) r'r(r)r*r staticmethodr>r#r+rrrr r.=s   r.cs,eZdZdZfddZfddZZS)MyGPT2LMHeadModelzSHere we wrap a class for Onnx model conversion for GPT2LMHeadModel with past state.cst|dSrrrrrr rdszMyGPT2LMHeadModel.__init__cs&tj||||dd}t||jjSr?rCrErrr r#gszMyGPT2LMHeadModel.forwardr&rrrr rHas rHcs,eZdZdZfddZfddZZS)MyGPT2LMHeadModel_NoPaddingaHere we wrap a class for Onnx model conversion for GPT2LMHeadModel with past state and no padding. When you always use batch_size=1 in inference, there is no padding in inputs. In such case, position_ids and attention_mask need no be in inputs. cst|dSrrrrrr rysz$MyGPT2LMHeadModel_NoPadding.__init__cs"tj||dd}t||jjS)NF)rBr"rC)rr%rFr;rrr r#|sz#MyGPT2LMHeadModel_NoPadding.forwardr&rrrr rIss rIlogitsTF last_state)rZGPT2LMHeadModel_NoPaddingrc@s8eZdZddZedddZedddZdd Zd S) Gpt2InputscCs||_||_||_||_dSr)r%r@rArF)rr%r@rArFrrr rszGpt2Inputs.__init__)returncCs0dd|j|j|jfD}|jr,||j|S)NcSsg|]}|dk r|qSrr.0vrrr sz&Gpt2Inputs.to_list..)r%r@rArFextend)r input_listrrr to_lists zGpt2Inputs.to_listcCs"tdd|j|j|j|jfDS)Ncss|]}|dk r|VqdSrrrNrrr sz&Gpt2Inputs.to_tuple..)r3r%r@rArF)rrrr to_tupleszGpt2Inputs.to_tuplecCsTd}|jdk r2|jjtjkr,|jjtjdn|j}dd|jD}t|j|j ||S)NdtypecSsg|]}|jtjdqS)rW)tor9float32rOprrr rQsz&Gpt2Inputs.to_fp32..) rArXr9float16rYrZrFrLr%r@)rrArFrrr to_fp32s  zGpt2Inputs.to_fp32N) r'r(r)rrrTrrVr^rrrr rLsrLc@s^eZdZdZedddejejejdfeeeeeeeeje e e ej ej ej e e dddZ edIeeee eeeeefdd d Zed d ZedJd dZedKddZedLddZedMddZeddddejejejfee e e e ej ej ej dddZedNddZedddd d!gfeeed"d#d$ZedOe ed%d&d'ZedPe ed%d(d)Zed*d+ZedQd,d-ZedRe eeejfeeeefee e d.d/d0Z ed1d2Z!ed3d4Z"edd5d5d6d7ddddejejejdddfd8d9Z#edd:ddddejejejd;d7dZ$edSd?d@Z%eddddAdBdCdDgfedEdFdGZ&dHS)T Gpt2HelperzEA helper class for Gpt2 model conversion, inference and verification.FT) batch_sizepast_sequence_lengthsequence_lengthnum_attention_heads hidden_sizer< vocab_sizedevicer]has_position_idshas_attention_maskinput_ids_dtypeposition_ids_dtypeattention_mask_dtypeleft_side_paddingrMcs$|r tjntjd|||t||gfddt|D}tjd|d||f| d}d}| r||}tj||g| d}|dkrt|D]>}td|d}|rd||d|f<qd||||df<qd}| r| d d}| |dkd|dd|df | }t ||||S) zCreate random inputs for GPT2 model. Returns torch tensors of input_ids, position_ids, attention_mask and a list of past state tensors. r0cs$g|]}tjdddqS)rXrfg@?)r9Zrand)rO_rfZ float_typeZ past_shaperr rQsz/Gpt2Helper.get_dummy_inputs..rr/)lowhighsizerXrfNrm) r9r]rZintr7randintZonesrandomlongZcumsumZ masked_fill_rYrL)r`rarbrcrdr<rerfr]rgrhrirjrkrlrFr%rAZtotal_sequence_lengthr=Zpadding_lengthr@rrpr get_dummy_inputssF  zGpt2Helper.get_dummy_inputsr)r`rarbr model_classrMcCs~|j}|j}|j}|j}t|d} ||| dkr4|n|g} d||||t||g} | | i} t|D]} | | dt| <qd| S)zAReturns a dictionary with output name as key, and shape as value.r/rJr0present_)rcrdZnum_hidden_layersre MODEL_CLASSESrur7str)r`rarbrrzrcrdr<reZ output_nameZlast_state_shapeZpresent_state_shape output_shapesr=rrr get_output_shapess&    zGpt2Helper.get_output_shapescCsZ|D]P}||kst||}t|||krtjt|||j|jd||<qdS)Nrm)r6numpyprodZnelementr9emptyrXrf)output_buffersr~keybufferrrr auto_increase_buffer_sizes  z$Gpt2Helper.auto_increase_buffer_sizecCsD|r tjntj}i}|D]"\}}tjt|||d||<q|S)zpReturns a dictionary of output name as key, and 1D tensor as value. The tensor has enough space for given shape.rm)r9r]rZitemsrrr)r~rf is_float16Z data_typernameshaperrr get_output_bufferss zGpt2Helper.get_output_bufferscCsL|d}t||d}|r>t|t|dSt|SdS)zGReturns the maximum difference between PyTorch and OnnxRuntime outputs.rư>N)cpurabsamax) torch_outputs ort_outputsrelativeZexpected_outputsdiffrrr diff_outputs&s zGpt2Helper.diff_outputsMbP?c Kstj|d|d||d}td||}t|d}t|D]R}tj|d||d|||d}td|d|d||o|}qJ|st||} t d| d |S) zReturns True if torch and ORT outputs are close for given thresholds, and False otherwise. Note: need kwargs since Gpt2BeamSearchHelper.compare_outputs has an extra parameter model_class r)rtolatolz9PyTorch and OnnxRuntime output 0 (last_state) are close: r/zPyTorch and OnnxRuntime layer z state (present_z ) are close:z@PyTorch and OnnxRuntime results are not all close: max_abs_diff=.5f) rallcloserloggerdebugr5r7r_rinfo) rrrrkwargsis_close is_all_closeZ num_layerslayer max_abs_diffrrr compare_outputs0s""     zGpt2Helper.compare_outputsrc Csd}d}g}g}tt|D]}||}|dkr:|dn|d|d} tj|| |dd} |tt| ||o| }t|  rt d|dt |  rt d|dt| rt d |dt | rt d |dt || } t| | j} |d | | d d | d || d dt| | d |dkrttj|dd|j} ttj| dd| j}t| |}q|t|}|t||||fS)aCompare outputs from PyTorch and OnnxRuntime Args: torch_outputs (Tuple[Torch.Tensor]): PyTorch model output ort_outputs (List[numpy.ndarray]): OnnxRuntime output atol (float, optional): Absolute tollerance. Defaults to 1e-06. Returns: is_all_close(bool): whether all elements are close. max_abs_diff(float): maximum absolute difference. messages(str): a list of debug message for each output TFrr/)rrzPyTorch output z has nanz has infz ORT output zdiff=z.9fz index=z ort=z torch=N)Zaxis)r7r5rrrr8rrisnananyrrisinffabsZ unravel_indexZargmaxrfloatZ array_equalindexmax)rrrris_top1_matchedZ max_diffsmessagesr=Z ort_outputZ torch_outputrridxZ ort_max_indexZtorch_max_indexmax_diff_output_indexrrr compare_outputs_v2KsD(0zGpt2Helper.compare_outputs_v2)onnx_model_pathverboseuse_external_data_formatrgrhrirjrkc Cs|j} | j} tjddd| j| j| | j|d||||| d} | } t || }W5QRXddt | D}ddt | D}|dj d| jks|dj d| jkst |dj d| jkrd nd f|}d d d d|dd d di}|D]}d dd||<q|D]}d dd||<qd g}|r@d d d|d<| d|r^d dd|d<| d||t|dkrt|d| kst td| jj d| jdj d|dj d|ddj t|jjddd|r^tj}tj|d}t|jjdddt|t| |d|||ddd|d tj|dd}t j!||dddW5QRXn"t|t| |d|||ddd|d dS) z1Export GPT-2 model with past state to ONNX model.r/F)r`rarbrcrdr<rerfr]rgrhrirjrkcSsg|]}d|qS)past_rrOr=rrr rQsz*Gpt2Helper.export_onnx..cSsg|]}d|qS)r{rrrrr rQsrr0rJrKr%r`Zseq_len)rr/Z past_seq_len)r/Z total_seq_lenr@rAzShapes: input_ids=z past=z output=z present=T)parentsexist_okz gpt2.onnx ) argsfZ export_params input_names output_names dynamic_axesZ opset_versionZdo_constant_foldingrr)Zload_external_data)Zsave_as_external_dataZall_tensors_to_one_fileN)"rrDr_ryrcrdrerTr9no_gradr7rr6r8rRr5rrr%rFrparentmkdirtempfileTemporaryDirectoryospathjoinr r3onnxZ load_modelr save)modelrfrrrrgrhrirjrkrr< dummy_inputsrSoutputsZ past_namesZ present_namesrrrrZ tmp_dir_nameZtemp_onnx_model_pathrrr export_onnxs ,"    $6 zGpt2Helper.export_onnxrc Ksftd} t|d||d| dd} |rV|r2t| n$d|krBd|d<| jfddi|| ||| S)zHOptimize ONNX model with an option to convert it to use mixed precision.rrF) model_typeZ num_headsrd opt_leveloptimization_optionsZuse_gpu keep_io_typesuse_symbolic_shape_inferT)r r r_auto_mixed_precisionconvert_float_to_float16Zsave_model_to_file) rZoptimized_model_pathrrcrdrrstagerrmrrr optimize_onnxs$    zGpt2Helper.optimize_onnxAddZLayerNormalizationZSkipLayerNormalizationZFastGeluZEmbedLayerNormalization) onnx_model op_block_listcCsTdd|D}t|}||}td|d||jdj}d}|}||ksbt ||}d} |j dkr|} td |jd} |j D]} | | } | dk rqqt | } td |jd | | d k}ntd |j d|jg} g}|s| dk r|g} | jg}| |||d}td||jfddi||S)a?Convert GPT-2 model to mixed precision. It detects whether original model has fp16 weights, and set parameters for float16 conversion automatically. Args: onnx_model (OnnxModel): optimized ONNX model op_block_list (List[str], optional): operators to compute in fp32. Defaults to ["Add", "LayerNormalization", "SkipLayerNormalization", "FastGelu", "EmbedLayerNormalization"] Returns: parameters(dict): a dictionary of parameters used in float16 conversion cSsh|] }|jqSr)op_type)rOnoderrr 1sz2Gpt2Helper.auto_mixed_precision..z fp32 op: z fp16 op: rFNZMatMulz#Found last MatMul node for logits: z3max diff of converting weights in last MatMul node : rz-Failed to find MatMul node for logits. Found z of node )rrnode_block_listZforce_fp16_initializersz!auto_mixed_precision parameters: rT)Znodesset differencerrgraphoutputroutput_name_to_noder6rinputZget_initializerrrwarningr)rrZ op_full_setZ fp32_op_setZ fp16_op_setZlogits_output_nameZis_weight_fp16_precisionrrZlast_matmul_nodeZ initializerrZmax_diffrr parametersrrr rsF      zGpt2Helper.auto_mixed_precision)inputs total_runsc Cstd|}t||}W5QRX|dkr>|Sg}t6t|D]&}t}||}|t|qTW5QRXt |dt |}td t |d||fS)zfRun inference of PyTorch model, and returns average latency in ms when total_runs > 0 besides outputs.zstart pytorch_inferencerzPyTorch inference time = {} ms.2f) rrr^rTr9rr7timer8sumr5format) rrrrSrlatencyrostartaverage_latencyrrr pytorch_inferencecs     zGpt2Helper.pytorch_inferencec Cs"tddt|ji}|jdk r\t|jD]$\}}t||d|<q6|jdk r~t|j|d<|j dk rt|j |d<| d|}|dkr|Sg}t |D]*}t } | d|}| t | qt|dt|} td t| d || fS) zcRun inference of ONNX model, and returns average latency in ms when total_runs > 0 besides outputs.zstart onnxruntime_inferencer%NrrAr@rrz"OnnxRuntime Inference time = {} msr)rrrZascontiguousarrayr%rrF enumeraterAr@runr7rr8rr5r) ort_sessionrrZ ort_inputsr=Zpast_irrrorrrrr onnxruntime_inference}s(       z Gpt2Helper.onnxruntime_inferencec Cst|||||||S)z)Returnas IO binding object for a session.)r prepare_io_binding)rr%r@rArFrr~rrr rs zGpt2Helper.prepare_io_bindingcCst||||S)z3Copy results to cpu. Returns a list of numpy array.)r "get_outputs_from_io_binding_buffer)rrr~ return_numpyrrr rs z-Gpt2Helper.get_outputs_from_io_binding_buffer)rrr~rrinclude_copy_output_latencyc Cstdt||j|j|j|j||}||t ||||}|dkrN|Sg} t |D]<} t } |||rt ||||} | t | qZt | dt| } td| || fS)zUInference with IO binding. Returns outputs, and optional latency when total_runs > 0.z*start onnxruntime_inference_with_binded_iorrz4OnnxRuntime with IO binding inference time = %.2f ms)rrr_rr%r@rArFZrun_with_iobindingrr7rr8rr5) rrrr~rrrZ io_bindingrrrorrrrr $onnxruntime_inference_with_binded_iosD     z/Gpt2Helper.onnxruntime_inference_with_binded_ioc Cs|td|dd}t||W5QRXtd|dtd|dd}t||W5QRXtd|ddS)NZ ort_outputs_.picklewbz$ORT output are saved to ort_outputs_Ztorch_outputs_z(Torch output are saved to torch_outputs_openpickledumprr)r=rrrrrr save_outputss zGpt2Helper.save_outputsc Cs@td|dd}t||W5QRXtd|ddS)NZ dummy_inputs_rrz!inputs are saved to dummy_inputs_r)r=rrrrrrr save_inputsszGpt2Helper.save_inputsri'r/c,s0|j}td|dd|d|d| d|d d}d }d }d }|rjt||||| }t|||}d }d }gd g|}|}t|D]}t|}t d |}|d krd n t d |}t d |} t d| d|dtj | |||j |j |j|j||| | | | |dd}!t||!}"|r:t||!}#n"t| |||| }$t||!||$}#tj|"|#|d\}%}&}'}(})t|&s|&|%r|d 7}|)r|d 7}||d 7<|r&|%s&td|d| d|d|d|& t|(D]0\}}*td|d||jd|*q|rt|&sB|&d|krt||!t||#|"qrxfdddD}+ndddD}+|d ||+d!<fd"d#|D|+d$<|d ||+d%<|td ||+d&<td'|d(|d)|td*||d+|kr,td,t|d|d-d.|+S)/zKGenerate random inputs and compare the results of PyTorch and Onnx Runtime.zRunning parity test (atol=z , test_cases=z, runs=z, use_io_binding=z, model_class=z , is_float16=z) ...r0Nrr/z#Running parity test for batch_size=z past_sequence_length=z...T)rirjrkrl)rz test_case=z batch_size=z sequence_length=z MaxDiff= z: Name=z, dcs$i|]}d|t|dqS)max_diff_percentile_r)rZ percentiler[)max_abs_diff_listrr psz*Gpt2Helper.test_parity..)2Z_ccSsi|]}d|dqS)rnanrr[rrr rtsrnZtop1_match_ratecsg|]}|dqS)rnr)rOx)test_cases_per_runrr rQwsz*Gpt2Helper.test_parity..Ztop1_match_rate_per_runZdiff_pass_rateZnan_ratezParity Test Cases=z ; Passed=z; Nan=z; Top1_Matched=gffffff?zParity is good: passed rate=z.0f%)rrrr_rrr7rurwrvrryrcrdrDrerrrrrrr8r get_outputsrrrr5),rrrfrrrrruse_io_bindingrzrgrhrirjrkrrZenable_pickle_outputrZmax_batch_sizeZmax_past_seq_lenZ max_seq_lenrZmax_output_shapesZpassed_test_casesZtop1_matched_casesZtop1_matched_cases_per_runZtotal_test_casesr=Zrun_idrbrar`rrrr~rrrrrmessager;r)rrr test_paritys(          *  " zGpt2Helper.test_parityrr cCs|j}d}|r.t| || ||}t|||}tj| || |j|j|j|j||||| | | d}|rtt |||\}}nt |||||\}}|S)zCGenerate random inputs and measure average latency of Onnx Runtime.N)rirjrk) rr_rrryrcrdrDrerr)rrrfrrrrzrgrhrirjrkr`rbrarrr~rrorrrr test_performancesJzGpt2Helper.test_performancec Cs:tjddd|j|j|j|j|d||d }tj ||S)zJIT trace for TorchScript.r/F) r`rarbrcrdr<rerfr]rgrh) r_ryrcrdrDrerTr9Zjittrace)rrrfrgrhrSrrr torchscripts zGpt2Helper.torchscriptrawfp32fp16int8)rzc Cs|}tj|r t|jd}n|dd|dkrB|d|7}|rN|d7}|rdddd d }d D]}tj||||} tj| rf||krzt | t d | Wn:t k r} zt d | d | j W5d} ~ XYnXqft d|d| qftjtj|||dtjtj||d|dtjtj||d|dtjtj||d |dd Stj||dtj||dtj||dtj||dd S)z=Build a path name for given model based on given attributes.rt/rroZ_pastZ_fp32Z_fp16Z_int8)rrrrzRemoved the existed directory: zFailed to remove the directory rNzDirectory for z existed: z.onnxz _fp32.onnxz _fp16.onnxz _int8.onnx)rrisdirrpartssplitrexistsshutilrmtreerrOSErrorstrerror) output_dirZmodel_name_or_pathrzZhas_pastZ new_folderZremove_existingZ model_namesuffixrZnew_direrrr get_onnx_pathssN     ,zGpt2Helper.get_onnx_pathsN)r)F)F)rr)r)FFr)r)r)T)rTF)TT)'r'r(r)r*rGr9Zint32rurfboolrXrLryrr}rrrrrrrrrrr rrrrrZTensorrrrrrrr)rrrr r_s& @"     5w #E   2    6  r_)4loggingrrrwr"rrpathlibrtypingrrrrrrr9Zbenchmark_helperrr]rZfusion_optionsr Zio_binding_helperr rr Z optimizerr Ztorch_onnx_export_helperr Z transformersrrrr getLoggerr'rZPRETRAINED_GPT2_MODELSZFLOAT32ZFLOAT16ZINT8ZDEFAULT_TOLERANCErr,r.rHrIr|rLr_rrrr sN           $