o  i @sddlmZddlmZddlmZddlmZmZm Z e dej ed\Z Z er/ddl mZn e dej ed\ZZ Gd d d Zd S) ) annotations) TYPE_CHECKING) trt_compile) IgniteInfo min_versionoptional_importz ignite.engineEvents)Enginer c@s.eZdZdZd ddZdd d Zdd d ZdS) TrtHandlera TrtHandler acts as an Ignite handler to apply TRT acceleration to the model. Usage example:: handler = TrtHandler(model=model, base_path="/test/checkpoint.pt", args={"precision": "fp16"}) handler.attach(engine) engine.run() NcCs||_||_||_||_dS)a Args: base_path: TRT path basename. TRT plan(s) saved to "base_path[.submodule].plan" args: passed to trt_compile(). See trt_compile() for details. submodule : Hierarchical ids of submodules to convert, e.g. 'image_decoder.decoder' N)model base_pathargs submodule)selfr r r rr\/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/handlers/trt_handler.py__init__#s zTrtHandler.__init__enginer returnNonecCs|j|_|tj|dS)g Args: engine: Ignite Engine, it can be a trainer, validator or evaluator. N)loggeradd_event_handlerrSTARTEDrrrrrattach/szTrtHandler.attachcCs t|j|j|j|j|jddS)r)r rrN)rr r r rrrrrr__call__7s zTrtHandler.__call__)NN)rr rr)__name__ __module__ __qualname____doc__rrrrrrrr s   r N) __future__rtypingrmonai.networksr monai.utilsrrrOPT_IMPORT_VERSIONr_ ignite.enginer r rrrrs