
    ^j                         d Z ddlZddlZddlZddlZddlmZ ddlmZ  e       Z	dde
deddfd	Zd
e
dej                  ddfdZde
dee
ef   fdZy)z?TensorRT export helpers: trtexec invocation and output parsing.    N)Any)
get_loggercommanddry_runreturnz subprocess.CompletedProcess[str]c                 `   |r0t         j                  dt        j                  dd       d|  d       	 t	        j
                  | dddd      }|S # t        j                  $ rJ}t         j                  d|j                          t         j                  d	|j                           d }~ww xY w)
Nz
CUDA_VISIBLE_DEVICES=CUDA_VISIBLE_DEVICES  
T)shellcapture_outputtextcheckzCommand failed with exit code zError output:
)
loggerinfoosgetenv
subprocessrunCalledProcessErrorerror
returncodestderr)r   r   resultes       b/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/rfdetr/export/_tensorrt.pyrun_command_shellr      s    -bii8NPR.S-TTUV]U^^`abtDt[_`(( 5all^DEqxxj12s   A B-#AB((B-onnx_dirargsc                 z   | j                  dd      }dj                  dd|  d| dd|j                  rd	nd
 g      }|j                  rD| j                  dd      }dj                  dd| dd|g      }t        j                  d|        n|}t        ||j                        }t        |j                         y )Nz.onnxz.enginer   trtexecz--onnx=z--saveEngine=z#--memPoolSize=workspace:4096 --fp16zF--useCudaGraph --useSpinWait --warmUp=500 --avgRuns=1000 --duration=10z	--verboser
   z	.nsys-repznsys profilez	--output=z--trace=cuda,nvtxz--force-overwrite truezProfile data will be saved to: )
replacejoinverboseprofiler   r   r   r   parse_trtexec_outputstdout)r   r    
engine_dirtrt_commandprofile_dirr   outputs          r   r"   r"   "   s    !!'95J ((hZ J<(1T"ll{3	
	K ||&&w<((y68KMegrs
 	5k]CDw5F'    output_textc           
         t         j                  |        d}d}d}d}d}i }t        j                  ||       x}rv|j	                  t        |j                  d            t        |j                  d            t        |j                  d            t        |j                  d	            d
       t        j                  ||       x}r]|j	                  t        |j                  d            t        |j                  d            t        |j                  d            d       t        j                  ||       x}r]|j	                  t        |j                  d            t        |j                  d            t        |j                  d            d       t        j                  ||       x}r]|j	                  t        |j                  d            t        |j                  d            t        |j                  d            d       t        j                  ||       x}rt        |j                  d            |d<   |S )NzhGPU Compute Time: min = (\d+\.\d+) ms, max = (\d+\.\d+) ms, mean = (\d+\.\d+) ms, median = (\d+\.\d+) msz\Host to Device Transfer Time: min = (\d+\.\d+) ms, max = (\d+\.\d+) ms, mean = (\d+\.\d+) msz\Device to Host Transfer Time: min = (\d+\.\d+) ms, max = (\d+\.\d+) ms, mean = (\d+\.\d+) mszGLatency: min = (\d+\.\d+) ms, max = (\d+\.\d+) ms, mean = (\d+\.\d+) mszThroughput: (\d+\.\d+) qps            )compute_min_mscompute_max_mscompute_mean_mscompute_median_ms)
h2d_min_ms
h2d_max_msh2d_mean_ms)
d2h_min_ms
d2h_max_msd2h_mean_ms)latency_min_mslatency_max_mslatency_mean_msthroughput_qps)r   r   researchupdatefloatgroup)r.   gpu_compute_patternh2d_patternd2h_patternlatency_patternthroughput_patternstatsmatchs           r   r'   r'   ?   s   
KK 	t  rKqK`O6E 		-{;;u;"'A"7"'A"7#(Q#8%*5;;q>%:		
 		+{33u3#EKKN3#EKKN3$U[[^4	
 		+{33u3#EKKN3#EKKN3$U[[^4	
 		/;77u7"'A"7"'A"7#(Q#8	
 		,k::u:"'A"7Lr-   )F)__doc__argparser   rB   r   typingr   rfdetr.utilities.loggerr   r   strboolr   	Namespacer"   dictr'    r-   r   <module>rW      sz    F  	 	   .		s 	T 	>` 	(c (!3!3 ( (:9c 9d38n 9r-   