Source code for pytorch_lightning.utilities.model_summary
# Copyright The PyTorch Lightning team.## Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License."""Utilities related to model weights summary."""importcontextlibimportloggingfromcollectionsimportOrderedDictfromtypingimportAny,Dict,List,Optional,Tuple,Unionimportnumpyasnpimporttorchimporttorch.nnasnnfromtorchimportTensorfromtorch.utils.hooksimportRemovableHandleimportpytorch_lightningasplfrompytorch_lightning.utilities.warningsimportWarningCachelog=logging.getLogger(__name__)warning_cache=WarningCache()PARAMETER_NUM_UNITS=[" ","K","M","B","T"]UNKNOWN_SIZE="?"
[docs]classLayerSummary:"""Summary class for a single layer in a :class:`~pytorch_lightning.core.lightning.LightningModule`. It collects the following information: - Type of the layer (e.g. Linear, BatchNorm1d, ...) - Input shape - Output shape - Number of parameters The input and output shapes are only known after the example input array was passed through the model. Example:: >>> model = torch.nn.Conv2d(3, 8, 3) >>> summary = LayerSummary(model) >>> summary.num_parameters 224 >>> summary.layer_type 'Conv2d' >>> output = model(torch.rand(1, 3, 5, 5)) >>> summary.in_size [1, 3, 5, 5] >>> summary.out_size [1, 8, 3, 3] Args: module: A module to summarize """def__init__(self,module:nn.Module)->None:super().__init__()self._module=moduleself._hook_handle=self._register_hook()self._in_size:Optional[Union[str,List]]=Noneself._out_size:Optional[Union[str,List]]=Nonedef__del__(self)->None:self.detach_hook()def_register_hook(self)->Optional[RemovableHandle]:"""Registers a hook on the module that computes the input- and output size(s) on the first forward pass. If the hook is called, it will remove itself from the from the module, meaning that recursive models will only record their input- and output shapes once. Registering hooks on :class:`~torch.jit.ScriptModule` is not supported. Return: A handle for the installed hook, or ``None`` if registering the hook is not possible. """defhook(_:nn.Module,inp:Any,out:Any)->None:iflen(inp)==1:inp=inp[0]self._in_size=parse_batch_shape(inp)self._out_size=parse_batch_shape(out)assertself._hook_handleisnotNoneself._hook_handle.remove()handle=Noneifnotisinstance(self._module,torch.jit.ScriptModule):handle=self._module.register_forward_hook(hook)returnhandle
[docs]defdetach_hook(self)->None:"""Removes the forward hook if it was not already removed in the forward pass. Will be called after the summary is created. """ifself._hook_handleisnotNone:self._hook_handle.remove()
@propertydefin_size(self)->Union[str,List]:returnself._in_sizeorUNKNOWN_SIZE@propertydefout_size(self)->Union[str,List]:returnself._out_sizeorUNKNOWN_SIZE@propertydeflayer_type(self)->str:"""Returns the class name of the module."""returnstr(self._module.__class__.__name__)@propertydefnum_parameters(self)->int:"""Returns the number of parameters in this module."""returnsum(np.prod(p.shape)ifnot_is_lazy_weight_tensor(p)else0forpinself._module.parameters())
[docs]classModelSummary:"""Generates a summary of all layers in a :class:`~pytorch_lightning.core.lightning.LightningModule`. Args: model: The model to summarize (also referred to as the root module). max_depth: Maximum depth of modules to show. Use -1 to show all modules or 0 to show no summary. Defaults to 1. The string representation of this summary prints a table with columns containing the name, type and number of parameters for each layer. The root module may also have an attribute ``example_input_array`` as shown in the example below. If present, the root module will be called with it as input to determine the intermediate input- and output shapes of all layers. Supported are tensors and nested lists and tuples of tensors. All other types of inputs will be skipped and show as `?` in the summary table. The summary will also display `?` for layers not used in the forward pass. Example:: >>> import pytorch_lightning as pl >>> class LitModel(pl.LightningModule): ... ... def __init__(self): ... super().__init__() ... self.net = nn.Sequential(nn.Linear(256, 512), nn.BatchNorm1d(512)) ... self.example_input_array = torch.zeros(10, 256) # optional ... ... def forward(self, x): ... return self.net(x) ... >>> model = LitModel() >>> ModelSummary(model, max_depth=1) # doctest: +NORMALIZE_WHITESPACE | Name | Type | Params | In sizes | Out sizes ------------------------------------------------------------ 0 | net | Sequential | 132 K | [10, 256] | [10, 512] ------------------------------------------------------------ 132 K Trainable params 0 Non-trainable params 132 K Total params 0.530 Total estimated model params size (MB) >>> ModelSummary(model, max_depth=-1) # doctest: +NORMALIZE_WHITESPACE | Name | Type | Params | In sizes | Out sizes -------------------------------------------------------------- 0 | net | Sequential | 132 K | [10, 256] | [10, 512] 1 | net.0 | Linear | 131 K | [10, 256] | [10, 512] 2 | net.1 | BatchNorm1d | 1.0 K | [10, 512] | [10, 512] -------------------------------------------------------------- 132 K Trainable params 0 Non-trainable params 132 K Total params 0.530 Total estimated model params size (MB) """def__init__(self,model:"pl.LightningModule",max_depth:int=1)->None:self._model=modelifnotisinstance(max_depth,int)ormax_depth<-1:raiseValueError(f"`max_depth` can be -1, 0 or > 0, got {max_depth}.")self._max_depth=max_depthself._layer_summary=self.summarize()# 1 byte -> 8 bits# TODO: how do we compute precision_megabytes in case of mixed precision?precision=self._model.precisionifisinstance(self._model.precision,int)else32self._precision_megabytes=(precision/8.0)*1e-6@propertydefnamed_modules(self)->List[Tuple[str,nn.Module]]:mods:List[Tuple[str,nn.Module]]ifself._max_depth==0:mods=[]elifself._max_depth==1:# the children are the top-level modulesmods=list(self._model.named_children())else:mods=self._model.named_modules()mods=list(mods)[1:]# do not include root module (LightningModule)returnmods@propertydeflayer_names(self)->List[str]:returnlist(self._layer_summary.keys())@propertydeflayer_types(self)->List[str]:return[layer.layer_typeforlayerinself._layer_summary.values()]@propertydefin_sizes(self)->List:return[layer.in_sizeforlayerinself._layer_summary.values()]@propertydefout_sizes(self)->List:return[layer.out_sizeforlayerinself._layer_summary.values()]@propertydefparam_nums(self)->List[int]:return[layer.num_parametersforlayerinself._layer_summary.values()]@propertydeftotal_parameters(self)->int:returnsum(p.numel()ifnot_is_lazy_weight_tensor(p)else0forpinself._model.parameters())@propertydeftrainable_parameters(self)->int:returnsum(p.numel()ifnot_is_lazy_weight_tensor(p)else0forpinself._model.parameters()ifp.requires_grad)@propertydefmodel_size(self)->float:# todo: seems it does not work with quantized models - it returns 0.0returnself.total_parameters*self._precision_megabytesdefsummarize(self)->Dict[str,LayerSummary]:summary=OrderedDict((name,LayerSummary(module))forname,moduleinself.named_modules)ifself._model.example_input_arrayisnotNone:self._forward_example_input()forlayerinsummary.values():layer.detach_hook()ifself._max_depth>=1:# remove summary entries with depth > max_depthforkin[kforkinsummaryifk.count(".")>=self._max_depth]:delsummary[k]returnsummarydef_forward_example_input(self)->None:"""Run the example input through each layer to get input- and output sizes."""model=self._modeltrainer=self._model.trainerinput_=model.example_input_arrayinput_=model._apply_batch_transfer_handler(input_)mode=model.trainingmodel.eval()forward_context=contextlib.nullcontext()iftrainerisNoneelsetrainer.precision_plugin.forward_context()withtorch.no_grad(),forward_context:# let the model hooks collect the input- and output shapesifisinstance(input_,(list,tuple)):model(*input_)elifisinstance(input_,dict):model(**input_)else:model(input_)model.train(mode)# restore mode of moduledef_get_summary_data(self)->List[Tuple[str,List[str]]]:"""Makes a summary listing with: Layer Name, Layer Type, Number of Parameters, Input Sizes, Output Sizes, Model Size """arrays=[(" ",list(map(str,range(len(self._layer_summary))))),("Name",self.layer_names),("Type",self.layer_types),("Params",list(map(get_human_readable_count,self.param_nums))),]ifself._model.example_input_arrayisnotNone:arrays.append(("In sizes",[str(x)forxinself.in_sizes]))arrays.append(("Out sizes",[str(x)forxinself.out_sizes]))returnarraysdef__str__(self)->str:arrays=self._get_summary_data()total_parameters=self.total_parameterstrainable_parameters=self.trainable_parametersmodel_size=self.model_sizereturn_format_summary_table(total_parameters,trainable_parameters,model_size,*arrays)def__repr__(self)->str:returnstr(self)
defparse_batch_shape(batch:Any)->Union[str,List]:ifhasattr(batch,"shape"):returnlist(batch.shape)ifisinstance(batch,(list,tuple)):shape=[parse_batch_shape(el)forelinbatch]returnshapereturnUNKNOWN_SIZEdef_format_summary_table(total_parameters:int,trainable_parameters:int,model_size:float,*cols:Tuple[str,List[str]],)->str:"""Takes in a number of arrays, each specifying a column in the summary table, and combines them all into one big string defining the summary table that are nicely formatted."""n_rows=len(cols[0][1])n_cols=1+len(cols)# Get formatting width of each columncol_widths=[]forcincols:col_width=max(len(str(a))forainc[1])ifn_rowselse0col_width=max(col_width,len(c[0]))# minimum length is header lengthcol_widths.append(col_width)# Formattings="{:<{}}"total_width=sum(col_widths)+3*n_colsheader=[s.format(c[0],l)forc,linzip(cols,col_widths)]# Summary = header + divider + Rest of tablesummary=" | ".join(header)+"\n"+"-"*total_widthforiinrange(n_rows):line=[]forc,linzip(cols,col_widths):line.append(s.format(str(c[1][i]),l))summary+="\n"+" | ".join(line)summary+="\n"+"-"*total_widthsummary+="\n"+s.format(get_human_readable_count(trainable_parameters),10)summary+="Trainable params"summary+="\n"+s.format(get_human_readable_count(total_parameters-trainable_parameters),10)summary+="Non-trainable params"summary+="\n"+s.format(get_human_readable_count(total_parameters),10)summary+="Total params"summary+="\n"+s.format(get_formatted_model_size(model_size),10)summary+="Total estimated model params size (MB)"returnsummarydefget_formatted_model_size(total_model_size:float)->str:returnf"{total_model_size:,.3f}"
[docs]defget_human_readable_count(number:int)->str:"""Abbreviates an integer number with K, M, B, T for thousands, millions, billions and trillions, respectively. Examples: >>> get_human_readable_count(123) '123 ' >>> get_human_readable_count(1234) # (one thousand) '1.2 K' >>> get_human_readable_count(2e6) # (two million) '2.0 M' >>> get_human_readable_count(3e9) # (three billion) '3.0 B' >>> get_human_readable_count(4e14) # (four hundred trillion) '400 T' >>> get_human_readable_count(5e15) # (more than trillion) '5,000 T' Args: number: a positive integer number Return: A string formatted according to the pattern described above. """assertnumber>=0labels=PARAMETER_NUM_UNITSnum_digits=int(np.floor(np.log10(number))+1ifnumber>0else1)num_groups=int(np.ceil(num_digits/3))num_groups=min(num_groups,len(labels))# don't abbreviate beyond trillionsshift=-3*(num_groups-1)number=number*(10**shift)index=num_groups-1ifindex<1ornumber>=100:returnf"{int(number):,d}{labels[index]}"returnf"{number:,.1f}{labels[index]}"
def_is_lazy_weight_tensor(p:Tensor)->bool:fromtorch.nn.parameterimportUninitializedParameterifisinstance(p,UninitializedParameter):warning_cache.warn("A layer with UninitializedParameter was found. ""Thus, the total number of parameters detected may be inaccurate.")returnTruereturnFalse
[docs]defsummarize(lightning_module:"pl.LightningModule",max_depth:int=1)->ModelSummary:"""Summarize the LightningModule specified by `lightning_module`. Args: lightning_module: `LightningModule` to summarize. max_depth: The maximum depth of layer nesting that the summary will include. A value of 0 turns the layer summary off. Default: 1. Return: The model summary object """returnModelSummary(lightning_module,max_depth=max_depth)
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.