U Ph?@sddlmZddlZddlmZmZddlmZddlZ ddl Z ddl m Z ddl m mZddlmZddlmZddlmZmZddlmZd d d d d gZddddd ZGdd d ZGdddZGdd d eZGdd d eZGdd d eZdS)) annotationsN)CallableSequence)cast) NdarrayTensor)ScaleIntensity) ensure_tuple pytorch_after)default_upsamplerCAMGradCAM GradCAMppModelWithHooksdefault_normalizerr)xreturncCsDddddd}t|tjr*sz8default_normalizer.._compute..r)axis)rnpstack)rrrr_compute(s z$default_normalizer.._compute)device) isinstancetorchTensor as_tensordetachcpunumpyr!)rr rrrrs   c@sjeZdZdZddddddddZd d Zd d Zd ddddZddddddZdddZ ddZ dS)rzy A model wrapper to run model forward/backward steps and storing some intermediate feature/gradient information. F nn.Modulezstr | Sequence[str]bool) nn_moduletarget_layer_namesregister_forwardregister_backwardcCs||_t||_i|_i|_d|_d|_||_||_g}| D]\}}||jkrTq@| ||jrt ddrd|j kr|j drd|j d<| ||n||||jr@|||q@|jrt|t|jkrtd|jddS)ag Args: nn_module: the model to be wrapped. target_layer_names: the names of the layer to cache. register_forward: whether to cache the forward pass output corresponding to `target_layer_names`. register_backward: whether to cache the backward pass output corresponding to `target_layer_names`. NinplaceFz._hookrrDrFrMrrLrr=bszModelWithHooks.backward_hookcsfdd}|S)Ncs|j<dSNr6)rJ_inputoutputrLrrrMksz*ModelWithHooks.forward_hook.._hookrrNrrLrr@iszModelWithHooks.forward_hookz&str | Callable[[nn.Module], nn.Module])layer_idrcCs\t|r||jSt|trH|jD] \}}||kr&ttj|Sq&td|ddS)z Args: layer_id: a layer name string or a callable. If it is a callable such as `lambda m: m.fc`, this method will return the module `self.model.fc`. Returns: a submodule from self.model. zCould not find r2N) callabler3r"strr9rnnModuleNotImplementedError)rDrSrFrGrrr get_layerps   zModelWithHooks.get_layerz torch.Tensorint)logitsr8rcCs|dd|fSrO)squeeze)rDr[r8rrr class_scoreszModelWithHooks.class_scoreNc  sjj}jj|f|}|dkr6|ddn|_d\}}jrbtfddjD}jr |t t j_ j j j|djD]&} | jkrtd| d| d qtfd djD}|rj|||fS) Nr/)NNc3s|]}j|VqdSrOrPrlayerrDrr sz*ModelWithHooks.__call__..) retain_graphzBackward hook for z& is not triggered; `requires_grad` of z should be `True`.c3s"|]}|jkrj|VqdSrOrIr_rarrrbs )r3trainingevalmaxr8r-tupler4r.r]rrZr7 zero_gradsumbackwardr5rBrCtrain) rDrr8rckwargsrkr[actigradr`rrar__call__s(     zModelWithHooks.__call__cCs|jSrO)r3rarrrget_wrapped_netszModelWithHooks.get_wrapped_net)FF)NF) __name__ __module__ __qualname____doc__rHr=r@rYr]rorprrrrr2s+ c@sTeZdZdZeedfdddddddd d Zdd dZdddZddZ ddZ dS)CAMBasez% Base class for CAM methods. Tr)rUrr*None)r+r4 upsamplerpostprocessingr.rcCs8|t|ts"t||d|d|_n||_||_||_dS)NT)r-r.)r"rr+rwrx)rDr+r4rwrxr.rrrrHs  zCAMBase.__init__r'r^cKs$|jtj|d|ifd|i|jS)a Computes the actual feature map size given `nn_module` and the target_layer name. Args: input_size: shape of the input tensor device: the device used to initialise the input tensor layer_idx: index of the target layer if there are multiple target layers. Defaults to -1. kwargs: any extra arguments to be passed on to the module as part of its `__call__`. Returns: shape of the actual feature map. r! layer_idx) compute_mapr#zerosshape)rD input_sizer!ryrlrrrfeature_map_sizes zCAMBase.feature_map_sizeNcCs tdS)a Compute the actual feature map with input tensor `x`. Args: x: input to `nn_module`. class_idx: index of the class to be visualized. Default to `None` (computing `class_idx` from `argmax`) layer_idx: index of the target layer if there are multiple target layers. Defaults to -1. Returns: activation maps (raw outputs without upsampling/post-processing.) NrX)rDrr8ryrrrrzs zCAMBase.compute_mapcCs&|jdd}|||}||S)N)r|rwrx)rDacti_maprZ img_spatialrrr_upsample_and_post_processsz"CAMBase._upsample_and_post_processcCs tdSrOrrarrrroszCAMBase.__call__)r'r^)Nr^) rqrrrsrtr rrHr~rzrrorrrrrus rucsLeZdZdZdeefdddddddfd d Zdd dZdddZZ S)r a Compute class activation map from the last fully-connected layers before the spatial pooling. This implementation is based on: Zhou et al., Learning Deep Features for Discriminative Localization. CVPR '16, https://arxiv.org/abs/1512.04150 Examples .. code-block:: python import torch # densenet 2d from monai.networks.nets import DenseNet121 from monai.visualize import CAM model_2d = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3) cam = CAM(nn_module=model_2d, target_layers="class_layers.relu", fc_layers="class_layers.out") result = cam(x=torch.rand((1, 1, 48, 64))) # resnet 2d from monai.networks.nets import seresnet50 from monai.visualize import CAM model_2d = seresnet50(spatial_dims=2, in_channels=3, num_classes=4) cam = CAM(nn_module=model_2d, target_layers="layer4", fc_layers="last_linear") result = cam(x=torch.rand((2, 3, 48, 64))) N.B.: To help select the target layer, it may be useful to list all layers: .. code-block:: python for name, _ in model.named_modules(): print(name) See Also: - :py:class:`monai.visualize.class_activation_maps.GradCAM` fcr)rUzstr | Callablerrv)r+r4 fc_layersrwrxrcs tj||||dd||_dS)a' Args: nn_module: the model to be visualized target_layers: name of the model layer to generate the feature map. fc_layers: a string or a callable used to get fully-connected weights to compute activation map from the target_layers (without pooling). and evaluate it at every spatial location. upsampler: An upsampling method to upsample the output image. Default is N dimensional linear (bilinear, trilinear, etc.) depending on num spatial dimensions of input. postprocessing: a callable that applies on the upsampled output image. Default is normalizing between min=1 and max=0 (i.e., largest input will become 0 and smallest input will become 1). F)r+r4rwrxr.N)superrHr)rDr+r4rrwrx __class__rrrHsz CAM.__init__Nr^c  s|j|f|\}}}||}|dkr2|dd}|j^}} } tj||| dddd}|j|jtjfdd|Dddtjfddt |Dddj|df| S) Nr/r^r)dimcsg|]}|dqS)).rr)ra)rrrr-sz#CAM.compute_map..cs$g|]\}}|||dfqS)r/r)rrb)rRrrr.sr) r+rfr|r#splitreshaperYrr enumerate) rDrr8ryrlr[rm_rcspatialr)rrRrrz%s  zCAM.compute_mapcKs|j|||f|}|||S)a Compute the activation map with upsampling and postprocessing. Args: x: input tensor, shape must be compatible with `nn_module`. class_idx: index of the class to be visualized. Default to argmax(logits) layer_idx: index of the target layer if there are multiple target layers. Defaults to -1. kwargs: any extra arguments to be passed on to the module as part of its `__call__`. Returns: activation maps rzr)rDrr8ryrlrrrrro1s z CAM.__call__)Nr^)Nr^) rqrrrsrtr rrHrzro __classcell__rrrrr s- c@s$eZdZdZd ddZd ddZdS) r a Computes Gradient-weighted Class Activation Mapping (Grad-CAM). This implementation is based on: Selvaraju et al., Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization, https://arxiv.org/abs/1610.02391 Examples .. code-block:: python import torch # densenet 2d from monai.networks.nets import DenseNet121 from monai.visualize import GradCAM model_2d = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3) cam = GradCAM(nn_module=model_2d, target_layers="class_layers.relu") result = cam(x=torch.rand((1, 1, 48, 64))) # resnet 2d from monai.networks.nets import seresnet50 from monai.visualize import GradCAM model_2d = seresnet50(spatial_dims=2, in_channels=3, num_classes=4) cam = GradCAM(nn_module=model_2d, target_layers="layer4") result = cam(x=torch.rand((2, 3, 48, 64))) N.B.: To help select the target layer, it may be useful to list all layers: .. code-block:: python for name, _ in model.named_modules(): print(name) See Also: - :py:class:`monai.visualize.class_activation_maps.CAM` NFr^cKs|j|f||d|\}}}||||}}|j^} } } || | ddj| | fdgt| } | |jddd} t| S)Nr8rcr^rr/Tkeepdim)r+r|viewmeanrAriFrelu)rDrr8rcryrlrrmrnrrrweightsrrrrrzls  ,zGradCAM.compute_mapcKs&|j|f|||d|}|||S)aD Compute the activation map with upsampling and postprocessing. Args: x: input tensor, shape must be compatible with `nn_module`. class_idx: index of the class to be visualized. Default to argmax(logits) layer_idx: index of the target layer if there are multiple target layers. Defaults to -1. retain_graph: whether to retain_graph for torch module backward call. kwargs: any extra arguments to be passed on to the module as part of its `__call__`. Returns: activation maps )r8rcryr)rDrr8ryrcrlrrrrrotszGradCAM.__call__)NFr^)Nr^F)rqrrrsrtrzrorrrrr Bs) c@seZdZdZdddZdS)r aW Computes Gradient-weighted Class Activation Mapping (Grad-CAM++). This implementation is based on: Chattopadhyay et al., Grad-CAM++: Improved Visual Explanations for Deep Convolutional Networks, https://arxiv.org/abs/1710.11063 See Also: - :py:class:`monai.visualize.class_activation_maps.GradCAM` NFr^cKs|j|f||d|\}}}||||}}|j^} } } |d} | d||d| | ddj| | fdgt| } t| dk| t | } | | d}t t tj|jj|}||| | ddj| | fdgt| }||jddd }t |S) Nrrr^r/rgHz>Tr)r+r|powmulrrirAr#where ones_likedivrrrr$r7exp)rDrr8rcryrlrrmrnrrrZalpha_nrZalpha_dralphaZ relu_gradrrrrrrzs  B0zGradCAMpp.compute_map)NFr^)rqrrrsrtrzrrrrr s ) __future__rrBcollections.abcrrtypingrr(rr#torch.nnrVtorch.nn.functional functionalr monai.configrmonai.transformsr monai.utilsrr Zmonai.visualize.visualizerr __all__rrrur r r rrrr s$      m>eD