o  i@sdZddlmZddlmZddlmZddlmZddl m Z ddl m Z ddl mZdd lmZdd lmZmZed d d \ZZdgZGdddZdS)zX Decorators and context managers for NVIDIA Tools Extension to profile MONAI components ) annotations) defaultdict)wraps)Any)Function)Module) Optimizer)Dataset) ensure_tupleoptional_importztorch._C._nvtxz:NVTX is not installed. Are you sure you have a CUDA build?) descriptorRangec@sbeZdZUdZeeZded<    dd ddZd!ddZ ddZ d"ddZ ddZ ddZ dS)#r a A decorator and context manager for NVIDIA Tools Extension (NVTX) Range for profiling. When used as a decorator it encloses a specific method of the object with an NVTX Range. When used as a context manager, it encloses the runtime context (created by with statement) with an NVTX Range. Args: name: the name to be associated to the range methods: (only when used as decorator) the name of a method (or a list of the name of the methods) to be wrapped by NVTX range. If None (default), the method(s) will be inferred based on the object's type for various MONAI components, such as Networks, Losses, Functions, Transforms, and Datasets. Otherwise, it look up predefined methods: "forward", "__call__", "__next__", "__getitem__" append_method_name: if append the name of the methods to be decorated to the range's name If None (default), it appends the method's name only if we are annotating more than one method. recursive: if set to True, it will recursively annotate every individual module in a list or in a chain of modules (chained using Compose). Default to False. dict name_counterNFname str | Nonemethodsstr | tuple[str, ...] | Noneappend_method_name bool | None recursiveboolreturnNonecCs||_||_||_||_dSN)rrrr)selfrrrrrR/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/utils/nvtx.py__init__7s zRange.__init__objrcCs8|jdur.t|ttfrt|dd|DSddlm}t||r+tdd|j|_d|_|j durgt|j }d|vs@d |vrH|d |j }|j |d 7<|j |d krd|d |j ||_ n||_ |j durs| ||_ nt|j |_ |jdurt|j d krd|_nd|_|j D] }||||jq|S) NTcss|] }tdd|VqdS)TrN)r ).0trrr Fsz!Range.__call__..r)Composer FZCuCIM TorchVision_)r isinstancelisttupletypeZmonai.transforms.composer$r transformsr__name__rr _get_methodr rlen_decorate_method)rrr$rmethodrrr__call__Cs2         zRange.__call__cs|r |jd|n|j|drt|}n|}t||tfdd}|drAGddd|}t|||||_dSt|||dS)N.__cs$t|i|}t|Sr)_nvtx rangePushArangePop)argskwargsoutputZ _temp_funcrrr range_wrapper}s z-Range._decorate_method..range_wrapperc@s eZdZdS)z7Range._decorate_method..NVTXRangeDecoratedClassN)r- __module__ __qualname__rrrrNVTXRangeDecoratedClasssr?)r startswithr+getattrrsetattr __class__)rrr1rownerr<r?rr;rr0os      zRange._decorate_methodr*cCst|tr dg}t |St|trdg}t |St|tr%ddg}t |St|tr1dg}t |Sgd}g}|D] }t||rE||q9t|dkrXtdt |d|t |S) Nforwardstepbackward __getitem__)rEr2__next__rHr'z*The method to be wrapped for this object [zm] is not recognized.The name of the method should be provided or the object should have one of these methods:) r(rrrr hasattrappendr/ ValueErrorr+r )rr method_listZdefault_methodsr1rrrr.s2         zRange._get_methodcCs>|jdur|jdd7<d|jd|_t|jdS)Ncontextr'Zcontext_)rrr5r6)rrrr __enter__s zRange.__enter__cCs tdSr)r5r7)rr+value tracebackrrr__exit__s zRange.__exit__)NNNF) rrrrrrrrrr)rrrr)rrrr*)r-r=r>__doc__rintr__annotations__rr2r0r.rOrRrrrrr !s  , ! N)rS __future__r collectionsr functoolsrtypingrZtorch.autogradrtorch.nnr torch.optimrtorch.utils.datar monai.utilsr r r5r&__all__r rrrrs