o $ ir@spdZddlmZddlZddlZddlZddlZddlmZddl m Z m Z m Z ddl mZddlmZddlmZmZddlZddlZddlmZdd lmZdd lmZdd lmZmZm Z dd l!m"Z"m#Z#dd l$m%Z%m&Z&e#d\Z'Z(e#d\Z)Z(e#d\Z*Z(e#d\Z+Z,e#dd\Z-Z(gdZ.ee/dZ0da1dddZ2ddZ3ddd!d"Z4dd#d$Z5ej6d%fdd.d/Z7ddd7d8Z8    ddd>d?Z9  dddCdDZ:dEejj;jfdLdMZ?ddQdRZ@ddSdTZAeddWdXZBeddYdZZCdd]d^ZD _   ` dddcddZEddhdiZF          j 0 ` ` kddddZG      j 0 ddddZHdddZI       0ddddZJddZKddZL ` `ddddZM ` `ddddZNe ` `ddddZOddddZPGdddejQZRejSejTfddZUejSejTfddZVGdddejjQZWGdddejjQZXdddZYdddZZdddZ[dddĄZ\dddƄZ]dS)zE Utilities and types for defining networks, these depend on PyTorch. ) annotationsN) OrderedDict)CallableMappingSequence)contextmanager)deepcopy)AnyIterable) get_logger)PathLike) ensure_tuplesave_objset_determinism)look_up_optionoptional_import)convert_to_dst_typeconvert_to_tensoronnxzonnx.reference onnxruntime polygraphytorch_tensorrtz1.4.0)one_hotpredict_segmentationnormalize_transformto_norm_affine CastTempType normal_init icnr_init pixelshufflepixelunshuffle eval_mode train_modeget_state_dictcopy_model_state save_stateconvert_to_onnxconvert_to_torchscriptconvert_to_trt meshgrid_ij meshgrid_xyreplace_modulesreplace_modules_templook_up_named_moduleset_named_modulehas_nvfuser_instance_normget_profile_shapes) module_name input_shape Sequence[int]dynamic_batchsizeSequence[int] | NonecCsPd dd}|r|||d}|||d}|||d }n|}}}|||fS) zb Given a sample input shape, calculate min/opt/max shapes according to dynamic_batchsize. r2r3 scale_numintcSsg|}||d<|S)Nr)r2r6Z scale_shaper8r8V/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/utils.pyscale_batch_sizeQsz,get_profile_shapes..scale_batch_sizerN)r2r3r6r7r8)r2r4r:min_input_shapeopt_input_shapemax_input_shaper8r8r9r0Ls   r0cCsVtdurtStddd\}atsdSz ddl}|dWtSty*daYtSw)zwhether the current environment has InstanceNorm3dNVFuser https://github.com/NVIDIA/apex/blob/23.05-devel/apex/normalization/instance_norm.py#L15-L16 Nzapex.normalizationZInstanceNorm3dNVFusernameFrZinstance_norm_nvfuser_cuda) _has_nvfuserr importlib import_module ImportError)_rCr8r8r9r/`s  r/FrAstrcCst|dd|Dd|d}|durdS|dkr|S|dD](}|r-|t|}q t|dd|Dddd}|durCdSt||}q |S) a get the named module in `mod` by the attribute name, for example ``look_up_named_module(net, "features.3.1.attn")`` Args: name: a string representing the module attribute. mod: a pytorch module to be searched (in ``mod.named_modules()``). print_all_options: whether to print all named modules when `name` is not found in `mod`. Defaults to False. Returns: the corresponding pytorch module's subcomponent such as ``net.features[3][1].attn`` cSh|]}|dqSrr8.0nr8r8r9 z'look_up_named_module..N)defaultprint_all_options.cSrHrIr8)rKitemr8r8r9rMrNF)r named_modulessplitisdigitr7getattr)rAmodrPname_strrLr8r8r9r-ts  r-cCsJ|dd}t|dkr|nd|f\}}|s|St||}t||||S)a look up `name` in `mod` and replace the layer with `new_layer`, return the updated `mod`. Args: mod: a pytorch module to be updated. name: a string representing the target module attribute. new_layer: a new module replacing the corresponding layer at ``mod.name``. Returns: an updated ``mod`` See also: :py:func:`monai.networks.utils.look_up_named_module`. rRr;r<rQ)rsplitlenr-setattr)rXrAZ new_layerZ mods_attrsubmodsattr_modr8r8r9r.s   r.r;labels torch.Tensor num_classesr7dtype torch.dtypedimreturncCs|j|dkrt|jdg|dt|j}t||}t|j}||dkr-td|||<tj|||jd}|j || dd}|S)a For every value v in `labels`, the value in the output will be either 1 or 0. Each vector along the `dim`-th dimension has the "one-hot" format, i.e., it has a total length of `num_classes`, with a one and `num_class-1` zeros. Note that this will include the background label, thus a binary mask should be treated as having two classes. Args: labels: input tensor of integers to be converted into the 'one-hot' format. Internally `labels` will be converted into integers `labels.long()`. num_classes: number of output channels, the corresponding length of `labels[dim]` will be converted to `num_classes` from `1`. dtype: the data type of the output one_hot label. dim: the dimension to be converted to `num_classes` channels from `1` channel, should be non-negative number. Example: For a tensor `labels` of dimensions [B]1[spatial_dims], return a tensor of dimensions `[B]N[spatial_dims]` when `num_classes=N` number of classes and `dim=1`. .. code-block:: python from monai.networks.utils import one_hot import torch a = torch.randint(0, 2, size=(1, 2, 2, 2)) out = one_hot(a, num_classes=2, dim=0) print(out.shape) # torch.Size([2, 2, 2, 2]) a = torch.randint(0, 2, size=(2, 1, 2, 2, 2)) out = one_hot(a, num_classes=2, dim=1) print(out.shape) # torch.Size([2, 2, 2, 2, 2]) r;z6labels should have a channel with length equal to one.)sizercdevice)reindexvalue) ndimlistshaper[torchreshapeAssertionErrorzerosrhscatter_long)r`rbrcrermshor8r8r9rs$"   rlogitsmutually_exclusivebool thresholdfloatr cCsB|s||kS|jddkrtd||kS|jdddS)a% Given the logits from a network, computing the segmentation by thresholding all values above 0 if multi-labels task, computing the `argmax` along the channel axis if multi-classes task, logits has shape `BCHW[D]`. Args: logits: raw data of model output. mutually_exclusive: if True, `logits` will be converted into a binary matrix using a combination of argmax, which is suitable for multi-classes task. Defaults to False. threshold: thresholding the prediction values if multi-labels task. r;zTsingle channel prediction, `mutually_exclusive=True` ignored, use threshold instead.T)keepdim)r7rmwarningswarnargmax)rwrxrzr8r8r9rs    rrhtorch.device | str | Nonetorch.dtype | None align_corners zero_centeredc Cst|tj|ddd}|jtj|d}|rDd||dk<d|r#|n|d}tt|tjdtj|df}|sCd|d d d f<n/d||d k<d|rQ|dn|}tt|tjdtj|df}|ssd|d|d d d f<| d j|d }d|_ |S)a Compute an affine matrix according to the input shape. The transform normalizes the homogeneous image coordinates to the range of `[-1, 1]`. Currently the following source coordinates are supported: - `align_corners=False`, `zero_centered=False`, normalizing from ``[-0.5, d-0.5]``. - `align_corners=True`, `zero_centered=False`, normalizing from ``[0, d-1]``. - `align_corners=False`, `zero_centered=True`, normalizing from ``[-(d-1)/2, (d-1)/2]``. - `align_corners=True`, `zero_centered=True`, normalizing from ``[-d/2, d/2]``. Args: shape: input spatial shape, a sequence of integers. device: device on which the returned affine will be allocated. dtype: data type of the returned affine align_corners: if True, consider -1 and 1 to refer to the centers of the corner pixels rather than the image corners. See also: https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample zero_centered: whether the coordinates are normalized from a zero-centered range, default to `False`. Setting this flag and `align_corners` will jointly specify the normalization source range. TF)rh wrap_sequence track_meta)rcrhg@?)r;gNrvrrc) rrnfloat64clonedetachtodiagcatones unsqueeze requires_grad)rmrhrcrrnormr8r8r9rs" " "raffinesrc_sizedst_sizecCst|tjstdt|jd|dks!|jd|jdkr,tdt |jd|jdd}|t |ks?|t |krQtd|dt |d t |d t ||j |j ||}t |d |j ||}||ttj||d d S)a Given ``affine`` defined for coordinates in the pixel space, compute the corresponding affine for the normalized coordinates. Args: affine: Nxdxd batched square matrix src_size: source image spatial shape dst_size: target image spatial shape align_corners: if True, consider -1 and 1 to refer to the centers of the corner pixels rather than the image corners. See also: https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample zero_centered: whether the coordinates are normalized from a zero-centered range, default to `False`. See also: :py:func:`monai.networks.utils.normalize_transform`. Raises: TypeError: When ``affine`` is not a ``torch.Tensor``. ValueError: When ``affine`` is not Nxdxd. ValueError: When ``src_size`` or ``dst_size`` dimensions differ from ``affine``. z%affine must be a torch.Tensor but is rRr;r<zaffine must be Nxdxd, got zaffine suggests z D, got src=zD, dst=zD.cpu)dstr) isinstancernTensor TypeErrortype__name__ ndimensionrm ValueErrortupler[rrhrcrnplinalginvnumpy)rrrrrsrZ src_xformZ dst_xformr8r8r9r!s  $$rg{Gz?std normal_func+Callable[[torch.Tensor, float, float], Any]NonecCs|jj}t|dddur7|ddks|ddkr7||jjd|t|dddur5tj|j jddSdS|ddkrQ||jjd |tj|j jd dSdS) a Initialize the weight and bias tensors of `m' and its submodules to values from a normal distribution with a stddev of `std'. Weight tensors of convolution and linear modules are initialized with a mean of 0, batch norm modules with a mean of 1. The callable `normal_func', used to assign values, should have the same arguments as its default normal_(). This can be used with `nn.Module.apply` to visit submodules of a network. weightNZConvrLinearrvbiasZ BatchNormrr) __class__rrWfindrdatanninit constant_r)mrrcnamer8r8r9rIs ,rc Cs|jj^}}}|t|}t||}t||g|}||}|dd}|||d}|dd|}|||g|}|dd}|jj |dS)z ICNR initialization for 2D/3D kernels adapted from Aitken et al.,2017 , "Checkerboard artifact free sub-pixel convolution". rr;rN) rrmr[r7rnrq transposerorepeatrcopy_) convZupsample_factorr out_channels in_channelsdims scale_factorZoc2kernelr8r8r9r^s    rx spatial_dimsrc s||}t|}|dd\}}|}||dkr.td|dd|d|d t||}||gfd d |ddD} ttddd|} | |d| d|} dd g} t|D] } | | | d|qd|||gg||dd}|| | }|S) a Apply pixel shuffle to the tensor `x` with spatial dimensions `spatial_dims` and scaling factor `scale_factor`. See: Shi et al., 2016, "Real-Time Single Image and Video Super-Resolution Using a nEfficient Sub-Pixel Convolutional Neural Network." See: Aitken et al., 2017, "Checkerboard artifact free sub-pixel convolution". Args: x: Input tensor with shape BCHW[D] spatial_dims: number of spatial dimensions, typically 2 or 3 for 2D or 3D scale_factor: factor to rescale the spatial dimensions by, must be >=1 Returns: Reshuffled version of `x`. Raises: ValueError: When input channels of `x` are not divisible by (scale_factor ** spatial_dims) Nr<rzNumber of input channels (z:) must be evenly divisible by scale_factor ** dimensions (z**=z).csg|]}|qSr8r8rKdfactorr8r9 rNz pixelshuffle..r;)rlrgrr7rangeextendropermute) rrrre input_size batch_sizechannelsZ scale_divisorZ org_channels output_sizeindicespermute_indicesidxr8rr9rrs0    " $rc s||}t|}|dd\}}|}||}tfdd|ddDr8tdd|dd||gfdd|ddD} ||gtfd d|ddDg} d d gd dt|Dd dt|D} || | }|| }|S)az Apply pixel unshuffle to the tensor `x` with spatial dimensions `spatial_dims` and scaling factor `scale_factor`. Inverse operation of pixelshuffle. See: Shi et al., 2016, "Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network." See: Aitken et al., 2017, "Checkerboard artifact free sub-pixel convolution". Args: x: Input tensor with shape BCHW[D] spatial_dims: number of spatial dimensions, typically 2 or 3 for 2D or 3D scale_factor: factor to reduce the spatial dimensions by, must be >=1 Returns: Unshuffled version of `x` with shape (B, C*(r**d), H/r, W/r) for 2D or (B, C*(r**d), D/r, H/r, W/r) for 3D, where r is the scale_factor and d is spatial_dims. Raises: ValueError: When spatial dimensions are not divisible by scale_factor Nr<c3s|] }|dkVqdS)rNr8rrr8r9 sz!pixelunshuffle..z3All spatial dimensions must be divisible by factor z. , spatial shape is: csg|]}|qSr8r8rrr8r9rrNz"pixelunshuffle..csg|]}|gqSr8r8rrr8r9rrr;cSsg|]}d|dqS)r<rr8rKir8r8r9rrcSsg|]}d|dqS)r<r8rr8r8r9rr)rlrganyrsumrror) rrrrerrrZscale_factor_multZ new_channelsrZ reshaped_sizerr8rr9r s  "(, r nets nn.Modulec gsdd|D}z-tdd|DVWdn1s wYW|D] }t|dr3|q(dS|D] }t|drC|q8w)a Set network(s) to eval mode and then return to original state at the end. Args: nets: Input network(s) Examples .. code-block:: python t=torch.rand(1,1,16,16) p=torch.nn.Conv2d(1,1,3) print(p.training) # True with eval_mode(p): print(p.training) # False print(p(t).sum().backward()) # will correctly raise an exception as gradients are calculated cSs g|] }t|dr|jr|qStraininghasattrrrJr8r8r9r zeval_mode..cS"g|] }t|dr |n|qS)eval)rrrJr8r8r9r"Ntrain)rnno_gradrr)rrrLr8r8r9r!s    r!c gsdd|D}z.tddd|DVWdn1s!wYW|D] }t|dr4|q)dS|D] }t|drD|q9w)a Set network(s) to train mode and then return to original state at the end. Args: nets: Input network(s) Examples .. code-block:: python t=torch.rand(1,1,16,16) p=torch.nn.Conv2d(1,1,3) p.eval() print(p.training) # False with train_mode(p): print(p.training) # True print(p(t).sum().backward()) # No exception cSs g|] }t|dr|js|qSrrrJr8r8r9rrztrain_mode..TcSr)r)rrrJr8r8r9rrNr)rnset_grad_enabledrr)rZ eval_listrLr8r8r9r"s    r"objtorch.nn.Module | MappingcCs0t|tjtjjfr |j}t|dr|S|S)z Get the state dict of input object if has `state_dict`, otherwise, return object directly. For data parallel model, automatically convert it to regular model first. Args: obj: input object to check and get the state_dict. state_dict)rr DataParallelparallelDistributedDataParallelmodulerr)rr8r8r9r#s r#rQTrsrccst|}tt|}fdd|D} t|t} } |D]#\} } || }||vrB|| vrB||j| jkrB| ||<| |q|rG|niD]9} ||| }||vr|| vr||j|| jkrwtd||jd|| jd|| ||<| |qI|dur|D]"\}}|||}|dur|d| vr|d||d<| |dqtt | } tt |  | }t d t | d t |d |rt|tjjrt|tjtjjfr|j}|||| |fS) a Compute a module state_dict, of which the keys are the same as `dst`. The values of `dst` are overwritten by the ones from `src` whenever their keys match. The method provides additional `dst_prefix` for the `dst` key when matching them. `mapping` can be a `{"src_key": "dst_key"}` dict, indicating `dst[dst_prefix + dst_key] = src[src_key]`. This function is mainly to return a model state dict for loading the `src` model state into the `dst` model, `src` and `dst` can have different dict keys, but their corresponding values normally have the same shape. Args: dst: a pytorch module or state dict to be updated. src: a pytorch module or state dict used to get the values used for the update. dst_prefix: `dst` key prefix, so that `dst[dst_prefix + src_key]` will be assigned to the value of `src[src_key]`. mapping: a `{"src_key": "dst_key"}` dict, indicating that `dst[dst_prefix + dst_key]` to be assigned to the value of `src[src_key]`. exclude_vars: a regular expression to match the `dst` variable names, so that their values are not overwritten by `src`. inplace: whether to set the `dst` module with the updated `state_dict` via `load_state_dict`. This option is only available when `dst` is a `torch.nn.Module`. filter_func: a filter function used to filter the weights to be loaded. See 'filter_swinunetr' in "monai.networks.nets.swin_unetr.py". Examples: .. code-block:: python from monai.networks.nets import BasicUNet from monai.networks.utils import copy_model_state model_a = BasicUNet(in_channels=1, out_channels=4) model_b = BasicUNet(in_channels=1, out_channels=2) model_a_b, changed, unchanged = copy_model_state( model_a, model_b, exclude_vars="conv_0.conv_0", inplace=False) # dst model updated: 76 of 82 variables. model_a.load_state_dict(model_a_b) # Returns: an OrderedDict of the updated `dst` state, the changed, and unchanged keys. c$h|]}rt|r|qSr8recompilesearchrKZs_key exclude_varsr8r9rMR$z#copy_model_state..zParam. shape changed from z to rRNrr;z'dst' model updated:  of z variables.)r#rrlitemsrmappendr}r~sortedset differenceloggerinfor[rrnrModulerrrrload_state_dict)rrZ dst_prefixmappingrinplace filter_funcsrc_dictZdst_dictZto_skipall_keysZ updated_keyssvalZdst_keykeyrjZnew_pairZunchanged_keysr8rr9r$sB1    $       r$torch.nn.Module | dictpathr cKsNi}t|tr|D] \}}t|||<q nt|}td||d|dS)a Save the state dict of input source data with PyTorch `save`. It can save `nn.Module`, `state_dict`, a dictionary of `nn.Module` or `state_dict`. And automatically convert the data parallel module to regular module. For example:: save_state(net, path) save_state(net.state_dict(), path) save_state({"net": net, "opt": opt}, path) net_dp = torch.nn.DataParallel(net) save_state(net_dp, path) Refer to: https://pytorch.org/ignite/v0.4.8/generated/ignite.handlers.DiskSaver.html. Args: src: input data to save, can be `nn.Module`, `state_dict`, a dictionary of `nn.Module` or `state_dict`. path: target file path to save the input object. kwargs: other args for the `save_obj` except for the `obj` and `path`. default `func` is `torch.save()`, details of the args: https://pytorch.org/docs/stable/generated/torch.save.html. )rrNr8)rdictrr#r)rrkwargsZckptkvr8r8r9r%ss r%-C6?modelinputs Sequence[Any] input_namesSequence[str] | None output_names opset_version int | None dynamic_axesDMapping[str, Mapping[int, str]] | Mapping[str, Sequence[int]] | Nonefilename Any | Noneverifytorch.device | Noneuse_ort ort_providerrtolatol use_tracedo_constant_foldingconstant_size_thresholdc! s|tmi}| r"|}|}d|vr!|dr!|r!||d<d}n tjj|fi|}t|s6t|tr:|f}nt|}d}|durLt }|j }n|}t d|tj j||f|||p`d|||d|t |}Wdn1sxwY|rtrddlm}m}|||d }||||rFt|trt|}durttjrd nd fd d |D}|}ttddt||d}Wdn1swYtdddd |jjD}tt|dd |D}| rtj |!| r| ndgd}|"d|}n t#$|}|"d|}tddt||D]\}} t|tj%rDtj&j'|(t)| |j*d| | dq(|S)a Utility to convert a model into ONNX model and optionally verify with ONNX or onnxruntime. See also: https://pytorch.org/docs/stable/onnx.html for how to convert a PyTorch model to ONNX. Args: model: source PyTorch model to save. inputs: input sample data used by pytorch.onnx.export. It is also used in ONNX model verification. input_names: optional input names of the ONNX model. output_names: optional output names of the ONNX model. opset_version: version of the (ai.onnx) opset to target. Must be >= 7 and not exceed the latest opset version supported by PyTorch, for more details: https://github.com/onnx/onnx/blob/main/docs/Operators.md and https://github.com/pytorch/pytorch/blob/master/torch/onnx/_constants.py dynamic_axes: specifies axes of tensors as dynamic (i.e. known only at run-time). If set to None, the exported model will have the shapes of all input and output tensors set to match given ones, for more details: https://pytorch.org/docs/stable/onnx.html#torch.onnx.export. filename: optional filename to save the ONNX model, if None, don't save the ONNX model. verify: whether to verify the ONNX model with ONNX or onnxruntime. device: target PyTorch device to verify the model, if None, use CUDA if available. use_ort: whether to use onnxruntime to verify the model. ort_provider": onnxruntime provider to use, default is ["CPUExecutionProvider"]. rtol: the relative tolerance when comparing the outputs of PyTorch model and TorchScript model. atol: the absolute tolerance when comparing the outputs of PyTorch model and TorchScript model. use_trace: whether to use `torch.jit.trace` to export the torchscript model. do_constant_folding: passed to onnx.export(). If True, extra polygraphy folding pass is done. constant_size_threshold: passed to polygrapy conatant forling, default = 16M kwargs: if use_trace=True: additional arguments to pass to torch.onnx.export() else: other arguments except `obj` for `torch.jit.script()` to convert model, for more details: https://pytorch.org/docs/master/generated/torch.jit.script.html. dynamorFNztorch_versioned_kwargs=)frrrrr!r)fold_constants save_onnx)Zsize_thresholdcudarc&g|]}t|tjr|n|qSr8rrnrrrrhr8r9r&z#convert_to_onnx..seedTcSsg|]}|jqSr8r@rr8r8r9rcSsg|]}|qSr8)rrrr8r8r9rrZCPUExecutionProvider)Z providersrrr)+rrnrjitscript is_tensorrrrtempfileNamedTemporaryFilerAprintrexportloadpolygraphy_importedZpolygraphy.backend.onnx.loaderr%r&rlvaluesrhr' is_availablerrr graphinputziprZInferenceSessionSerializeToStringrun onnxreferenceZReferenceEvaluatorrtesting assert_closerrrc)!rrrrrrrrrhrrrrr r!r"r Ztorch_versioned_kwargsZmode_to_exportZ onnx_inputs temp_filer$ onnx_modelr%r& torch_outZmodel_input_names input_dictZort_sessZonnx_outZsessr1r2r8r*r9r&s2   $           "r&filename_or_obj extra_files dict | NoneSequence[Any] | Nonec  s|t6|r |durtdtjj|fd|i| } n tjj|fi| } |dur7tjj| ||dWdn1sAwY|rЈdurXttj rUdnd|dur`tdfdd |D}|durstj |n| } | | }t#t d d t||} t d d t| |} t dd Wdn1swYt| | D]\}}t|tjst|tjrtjj||||d q| S) a Utility to convert a model into TorchScript model and save to file, with optional input / output data verification. Args: model: source PyTorch model to save. filename_or_obj: if not None, specify a file-like object (has to implement write and flush) or a string containing a file path name to save the TorchScript model. extra_files: map from filename to contents which will be stored as part of the save model file. for more details: https://pytorch.org/docs/stable/generated/torch.jit.save.html. verify: whether to verify the input and output of TorchScript model. if `filename_or_obj` is not None, load the saved TorchScript model and verify. inputs: input test data to verify model, should be a sequence of data, every item maps to a argument of `model()` function. device: target device to verify the model, if None, use CUDA if available. rtol: the relative tolerance when comparing the outputs of PyTorch model and TorchScript model. atol: the absolute tolerance when comparing the outputs of PyTorch model and TorchScript model. use_trace: whether to use `torch.jit.trace` to export the TorchScript model. kwargs: other arguments except `obj` for `torch.jit.script()` or `torch.jit.trace()` (if use_trace is True) to convert model, for more details: https://pytorch.org/docs/master/generated/torch.jit.script.html. Nz'Missing input data for tracing convert.example_inputs)rr$ _extra_filesr'r$Missing input data for verification.cr(r8r)rr*r8r9rGr+z*convert_to_torchscript..rr,r/)rrnrrr0tracer1saverhr'r:r7rrr r=rrrArB)rrIrJrrrhrrr r  script_moduleZts_modelrEZtorchscript_outrGrHr8r*r9r's@"         r' min_shape opt_shape max_shape precisioncCsPtdd\}} |||f} |sgn|}|sgn|}tj|||jj} || } | dt|j j >} | }|rF|j |dg| R| | | }||}|sqd}t|jD] }|||d7}q\td|| }|||dkr||jj| | |}t}||tjj |!t"d |||d }|S) af This function takes an ONNX model as input, exports it to a TensorRT engine, wraps the TensorRT engine to a TensorRT engine-based TorchScript model and return the TorchScript model. Args: onnx_model: the source ONNX model to compile. min_shape: the minimum input shape of the converted TensorRT model. opt_shape: the optimization input shape of the model, on which the TensorRT optimizes. max_shape: the maximum input shape of the converted TensorRT model. device: the target GPU index to convert and verify the model. precision: the weight precision of the converted TensorRT engine-based TorchScript model. Should be 'fp32' or 'fp16'. input_names: optional input names of the ONNX model. Should be a sequence like `['input_0', 'input_1', ..., 'input_N']` where N equals to the number of the model inputs. output_names: optional output names of the ONNX model. Should be a sequence like `['output_0', 'output_1', ..., 'output_N']` where N equals to the number of the model outputs. tensorrtz8.5.3r;rrQ z.TensorRT cannot parse the ONNX model, due to: fp16cuda:)rhZinput_binding_namesZoutput_binding_names)#rrnr' set_deviceLoggerWARNINGBuilderZcreate_networkr7ZNetworkDefinitionCreationFlagZEXPLICIT_BATCHZcreate_optimization_profile set_shapeZ OnnxParserparser>rZ num_errorsZ get_errordesc ExceptionZcreate_builder_configZadd_optimization_profileZset_flagZ BuilderFlagFP16Zbuild_serialized_networkioBytesIOwritertsZembed_engine_in_new_modulegetvalueDevice)rDrSrTrUrhrVrrtrtrF input_shapesrbuildernetworkprofileparsersuccessZparser_error_messagerconfigZserialized_enginer$ trt_modelr8r8r9_onnx_trt_compileZs@         rsZinput_0Zoutput_0{Gz?use_onnx bool | Noneonnx_input_namesonnx_output_namesc  Kstjs td|std|std|d|dur-t|dkr-td|d|r1|nd }td |}|d krBtj ntj }t t | |g}| |}t||\}}}|r| rjd d | Dni}|| rwdd | Dnit||| | ||d}t||||||| | d}nYt||||d}|| |t<tjj|d%tj|||dg}tj|f||td |dd| }Wdn1swYWdn1swY|rT|durtd|durtj|n|}t#td dt ||}td dt ||}tddWdn 1s,wYt||D]\}}t|tjsHt|tjrRtjj ||| | dq6|S)a Utility to export a model into a TensorRT engine-based TorchScript model with optional input / output data verification. There are two ways to export a model: 1, Torch-TensorRT way: PyTorch module ---> TorchScript module ---> TensorRT engine-based TorchScript. 2, ONNX-TensorRT way: PyTorch module ---> TorchScript module ---> ONNX model ---> TensorRT engine ---> TensorRT engine-based TorchScript. When exporting through the first way, some models suffer from the slowdown problem, since Torch-TensorRT may only convert a little part of the PyTorch model to the TensorRT engine. However when exporting through the second way, some Python data structures like `dict` are not supported. And some TorchScript models are not supported by the ONNX if exported through `torch.jit.script`. Args: model: a source PyTorch model to convert. precision: the weight precision of the converted TensorRT engine based TorchScript models. Should be 'fp32' or 'fp16'. input_shape: the input shape that is used to convert the model. Should be a list like [N, C, H, W] or [N, C, H, W, D]. dynamic_batchsize: a sequence with three elements to define the batch size range of the input for the model to be converted. Should be a sequence like [MIN_BATCH, OPT_BATCH, MAX_BATCH]. After converted, the batchsize of model input should between `MIN_BATCH` and `MAX_BATCH` and the `OPT_BATCH` is the best performance batchsize that the TensorRT tries to fit. The `OPT_BATCH` should be the most frequently used input batchsize in the application, default to None. use_trace: whether using `torch.jit.trace` to convert the PyTorch model to a TorchScript model and then convert to a TensorRT engine based TorchScript model or an ONNX model (if `use_onnx` is True), default to False. filename_or_obj: if not None, specify a file-like object (has to implement write and flush) or a string containing a file path name to load the TensorRT engine based TorchScript model for verifying. verify: whether to verify the input and output of the TensorRT engine based TorchScript model. device: the target GPU index to convert and verify the model. If None, use #0 GPU. use_onnx: whether to use the ONNX-TensorRT way to export the TensorRT engine-based TorchScript model. onnx_input_names: optional input names of the ONNX model. This arg is only useful when `use_onnx` is True. Should be a sequence like `('input_0', 'input_1', ..., 'input_N')` where N equals to the number of the model inputs. If not given, will use `('input_0',)`, which supposes the model only has one input. onnx_output_names: optional output names of the ONNX model. This arg is only useful when `use_onnx` is True. Should be a sequence like `('output_0', 'output_1', ..., 'output_N')` where N equals to the number of the model outputs. If not given, will use `('output_0',)`, which supposes the model only has one output. rtol: the relative tolerance when comparing the outputs between the PyTorch model and TensorRT model. atol: the absolute tolerance when comparing the outputs between the PyTorch model and TensorRT model. kwargs: other arguments except `module`, `inputs`, `enabled_precisions` and `device` for `torch_tensorrt.compile()` to compile model, for more details: https://pytorch.org/TensorRT/py_api/torch_tensorrt.html#torch-tensorrt-py. zCannot find any GPU devices.z*Missing the input shape for model convert.z@There is no dynamic batch range. The converted model only takes z shape input.NrzAThe dynamic batch range sequence should have 3 elements, but got z elements.rrZfp32cSi|]}|ddiqSrZ batchsizer8rKr r8r8r9 z"convert_to_trt..cSr|r}r8r~r8r8r9rr)r r)rSrTrUrhrVrr)rhrr r*)rSrTrUZ torchscript)renabled_precisionsrhirrOr,r/)!rnr'r:rbrr}r~r[rhfloat32halfrandr rrr0updater&rsr'rrInputrrir0r7rr=rrrArB)rrVr2r4r rIrrhrwryrzrrr  target_deviceZconvert_precisionrr=r>r?rZir_modelrrZinput_placeholderrEZtrt_outrGrHr8r8r9r(s :           r(cGs2tjjdurdtjjvrtj|ddiStj|S)Nindexingijrnmeshgrid__kwdefaults__tensorsr8r8r9r)/s r)cGsLtjjdurdtjjvrtj|ddiStj|d|dg|ddRS)Nrxyr;rr<rrr8r8r9r*6s$r*parenttorch.nn.Module new_moduleout!list[tuple[str, torch.nn.Module]] strict_match match_devicec s|rtdd|D}t|dkr||d|d}|dkrL|d|t|}||dd}g}t|||||fdd |D7}dS|rbt||} t|||||| fg7}dS|D]\} } || vryt|| t ||d d qfdS) zO Helper function for :py:class:`monai.networks.utils.replace_modules`. cSsh|]}|jqSr8r*rr8r8r9rMIr.z#_replace_modules..r;rrRrNcs&g|]}d|d|dfqS)rRrr;r8)rKr parent_namer8r9rVr+z$_replace_modules..T)r) rl parametersr[rrrW_replace_modulesr\rTr) rrArrrrdevicesr_out old_modulemod_namerFr8rr9r=s*       rcCsg}t|||||||S)a Replace sub-module(s) in a parent module. The name of the module to be replace can be nested e.g., `features.denseblock1.denselayer1.layers.relu1`. If this is the case (there are "." in the module name), then this function will recursively call itself. Args: parent: module that contains the module to be replaced name: name of module to be replaced. Can include ".". new_module: `torch.nn.Module` to be placed at position `name` inside `parent`. This will be deep copied if `strict_match == False` multiple instances are independent. strict_match: if `True`, module name must `== name`. If false then `name in named_modules()` will be used. `True` can be used to change just one module, whereas `False` can be used to replace all modules with similar name (e.g., `relu`). match_device: if `True`, the device of the new module will match the model. Requires all of `parent` to be on the same device. Returns: List of tuples of replaced modules. Element 0 is module name, element 1 is the replaced module. Raises: AttributeError: if `strict_match` is `True` and `name` is not a named module in `parent`. r)rrArrrrr8r8r9r+cs r+ccslg}z t||||||dVW|D]\}}t|||gd|dqdS|D]\}}t|||gd|dq&w)z Temporarily replace sub-module(s) in a parent module (context manager). See :py:class:`monai.networks.utils.replace_modules`. NT)rrr)rrArrrreplacedrr8r8r9r,s  r,c sdur dur tdt|}t}durEfdd|D}|D] \}}||vr5d|_||q$|jsDd|_td|dq$durwfd d|D}|D] \}}||vrgd|_||qV|jsvd|_td |dqVt t |d t |d dS) a  A utilty function to help freeze specific layers. Args: model: a source PyTorch model to freeze layer. freeze_vars: a regular expression to match the `model` variable names, so that their `requires_grad` will set to `False`. exclude_vars: a regular expression to match the `model` variable names, except for matched variable names, other `requires_grad` will set to `False`. Raises: ValueError: when freeze_vars and exclude_vars are both specified. NzEIncompatible values: freeze_vars and exclude_vars are both specified.crr8rr) freeze_varsr8r9rMrz freeze_layers..FTz!The freeze_vars does not include z0, but requires_grad is False, change it to True.crr8rrrr8r9rMrzThe exclude_vars includes rz variables frozen.) rr#rlnamed_parametersrrr}r~rrr[) rrrrZ frozen_keysZ to_freezerAparamZ to_excluder8)rrr9 freeze_layerss6   "rc(eZdZdZfddZddZZS)rz} Cast the input tensor to a temporary type before applying the submodule, and then cast it back to the initial type. cs t||_||_||_dSN)super__init__ initial_typetemporary_type submodule)selfrrrrr8r9rs  zCastTempType.__init__cCs@|j}||jkr||j}||}||jkr||j}|Sr)rcrrrr)rrrcr8r8r9forwards     zCastTempType.forwardr __module__ __qualname____doc__rr __classcell__r8r8rr9rs rcCs|j|kr |j|dS|S)zN Utility function to cast a single tensor from from_dtype to to_dtype r)rcr)r from_dtypeto_dtyper8r8r9 cast_tensorsrcsvt|tjr t|dSt|tr(i}|D] }t||d||<q|St|tr9tfdd|DSdS)zU Utility function to cast all tensors in a tuple from from_dtype to to_dtype rrc3s|] }t|dVqdS)rN)cast_all)rKyrr8r9rszcast_all..N)rrnrrrkeysrr)rrrnew_dictr r8rr9rs    rcr) CastToFloatzo Class used to add autocast protection for ONNX export for forward methods with single return vaue ct||_dSrrrrXrrXrr8r9r  zCastToFloat.__init__cCsT|j}tjddd|j|tj|}Wd|S1s#wY|S)Nr'Fenabled)rcrnautocastrXrrr)rrrcretr8r8r9rs zCastToFloat.forwardrr8r8rr9r rcr)CastToFloatAllzs Class used to add autocast protection for ONNX export for forward methods with multiple return values crrrrrr8r9rrzCastToFloatAll.__init__cGs`|dj}tjddd|jjt||tjd}Wdn1s#wYt|tj|dS)Nrr'Frr)rcrnrrXrrr)rargsrrr8r8r9rs zCastToFloatAll.forwardrr8r8rr9r rrbase_ttype[nn.Module]dest_t'Callable[[nn.Module], nn.Module | None]csdfdd }|S) z Generic function generator to replace base_t module with dest_t wrapper. Args: base_t : module type to replace dest_t : destination module type Returns: swap function to replace base_t module with dest_t rXrrfnn.Module | Nonecs |}|Srr8)rXrrr8r9 expansion_fn$sz!wrap_module..expansion_fnNrXrrfrr8rrrr8rr9 wrap_modules rcsdfdd }|S) a7 Generic function generator to replace base_t module with dest_t. base_t and dest_t should have same atrributes. No weights are copied. Args: base_t : module type to replace dest_t : destination module type Returns: swap function to replace base_t module with dest_t rXrrfrcs2tsdSj}fdd|D}|}|S)Ncsg|]}t|dqSr)rW)rKrArXr8r9r:rz8simple_replace..expansion_fn..)r __constants__)rX constantsrrrrrr9r6s z$simple_replace..expansion_fnNrr8rr8rr9simple_replace+s rrdict[str, nn.Module]cCs^|D](\}}|d}|}|ddD]}|j|}|dur"n|}q||j|d<q|S)a This function swaps nested modules as specified by "dot paths" in mod with a desired replacement. This allows for swapping nested modules through arbitrary levels if children NOTE: This occurs in place, if you want to preserve model then make sure to copy it first. rRNr)rrU_modules)rrrnew_modZ expanded_pathZ parent_modZsub_pathsubmodr8r8r9 _swap_modulesAs  r expansions2dict[str, Callable[[nn.Module], nn.Module | None]]cCsbi}|D]\}}t|j}||vr|||}|r|||<qtdt|dt|||S)a Top-level function to replace modules in model, specified by class name with a desired replacement. NOTE: This occurs in place, if you want to preserve model then make sure to copy it first. Args: model : top level module expansions : replacement dictionary: module class name -> replacement function generator Returns: model, possibly modified in-place zSwapped z modules)rTrrr5r[r)rrrrArZm_typeswappedr8r8r9replace_modules_by_typeWs    rcCsXtdttjtttjtttjtttjtttjtttjtd}t |||S)a5 Top-level function to add cast wrappers around modules known to cause issues for FP16/autocast ONNX export NOTE: This occurs in place, if you want to preserve model then make sure to copy it first. Args: model : top level module Returns: model, possibly modified in-place zAdding casts around norms...) BatchNorm1d BatchNorm2d BatchNorm3d LayerNormInstanceNorm1dInstanceNorm3d) r5rrrrrrrrr)rZcast_replacementsr8r8r9add_casts_around_normsqs       r)r2r3r4r5)F)rArG) r`rarbr7rcrdrer7rfra)Frv)rwrarxryrzr{rfr )NNFF) rhrrcrrryrryrfra)FF) rrarr3rr3rryrryrfra)rr{rrrfr)rrarr7rr7rfra)rr)rr)rQNNTN)rrrr)rrrr )NNNNNFNFNr rvTTr ) rrrrrrrrrrrrrrrryrhrrryrrrr{rr{r ryr!ryr"r7)NNFNNr rvF)rrrIrrJrKrryrrLrhrrr{rr{r ry)rSr3rTr3rUr3rhr7rVrGrrrr) NFNFNFrtrurvrv)rrrVrGr2r3r4r5r ryrIrrryrhrrwrxryrrzrrr{rr{)TT)rrrArGrrrrrryrryrfr) rrrArGrrrryrryrfr) rrrArGrrrryrry)NN)rr)rrrrrfr)rrrrrfr)rrrrrfr)rrrfr)^r __future__rrdrr3r} collectionsrcollections.abcrrr contextlibrcopyrtypingr r rrrntorch.nnrmonai.apps.utilsr monai.configr monai.utils.miscr rrmonai.utils.modulerrmonai.utils.type_conversionrrrrFr@rrr8r__all__rrrBr0r/r-r.r{rrrrrnormal_rkaiming_normal_rrr r!r"r#r$r%r&r'rsr(r)r*rr+r,rrrfloat16rrrrrrrrrrr8r8r8r9s                5 2 )   -* "  # U%  FN    * % ,