
    NgR              "           d dl mZmZmZ d dlZd Z	 	 	 	 	 	 	 	 	 	 	 	 	 	 ddej        j        ded	eej	                 d
e
de
de
de
dedeeeef         dee         de
de
dee
         de
dee         dee         f dZdS )    )OptionalTupleListNc                    dd l }|                                }|                    | |          }|                                d         j        }|                    g ||                                i          }|d         }|S )Nr   )onnxruntimeSessionOptionsInferenceSession
get_inputsnamerunnumpy)	onnx_fileexample_inputr   sess_optionssession
input_nameoutputs          K/var/www/html/ai-engine/env/lib/python3.11/site-packages/timm/utils/onnx.pyonnx_forwardr      s    --//L**9lCCG##%%a(-J[[j-*=*=*?*?@AAFAYFM    FT@   modeloutput_filer   trainingverbosecheckcheck_forward
batch_size
input_sizeopsetdynamic_sizeaten_fallbackkeep_initializers
use_dynamoinput_namesoutput_namesc                 f   dd l }|r+t          j         j        j        }|                                  n*t          j         j        j        }|                                  |H|s,t          | d          sJ | j        	                    d          }t          j
        |f|z   |          }t          j                    5   | |          }d d d            n# 1 swxY w Y   |pdg}|pdg}ddiddid}|
rd	|d         d
<   d|d         d<   |rt          j         j        j        }nt          j         j        j        }|rZt          j                             |
          }t          j                             | ||          }|                    |           d }n+t          j                             | |||d||||||	|          }|r|                    |          }|j                            |d           |r|sdd l}t/          ||          }|r|j                            |                                |d           |j                            |                                |                                d           d S |j                            |                                |d           d S d S d S d S )Nr   default_cfgr   )requires_gradinput0output0batch)r*   r+   height   width   )dynamic_shapes)export_optionsT)	r   export_paramsr   r%   r&   keep_initializers_as_inputsdynamic_axesopset_versionoperator_export_type)
full_check)decimal   )onnxtorchTrainingModeTRAININGtrainEVALevalhasattrr(   getrandnno_gradOperatorExportTypesONNX_ATEN_FALLBACKONNXExportOptionsdynamo_exportsave_exportloadcheckercheck_modelr   r   testingassert_almost_equal)r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r$   r%   r&   r;   training_modeoriginal_outr5   export_typer2   export_output	torch_out
onnx_modelnponnx_outs                              r   onnx_exportrZ      sL   $ KKK 
/8
/4

 	=5-00000*..|<<JZMJ$>hWWW 
 , ,u]++, , , , , , , , , , , , , , , +(K.I;L 'l7|DDL ,$,Xq!$+Xq! :j4Gj49 
111NN
00) 1 
 

 	;'''		J&&"#%(9%!, ' 
 
	  
ZYY{++
   === 	Z 	Z#K??H$
..y/@/@(TU.VVV
..|/A/A/C/CY__EVEV`a.bbbbb
..|/A/A/C/CXWX.YYYYY
Z 
Z	Z 	Z 	Z 	Zs   9CCC)NFFTFr   NNFFNFNN)typingr   r   r   r<   r   nnModulestrTensorboolintrZ    r   r   <module>rc      sz   ( ( ( ( ( ( ( ( ( (    15#+/#"#,0 !%"&![Z [Zx[Z[Z  -[Z 	[Z
 [Z [Z [Z [Z #sC-([Z }[Z [Z [Z $D>[Z [Z #Y[Z  3i![Z [Z [Z [Z [Z [Zr   