Source code for nemo.core.classes.module
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.## 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.fromcontextlibimportcontextmanagerimporttorchfromtorch.nnimportModulefromnemo.core.classes.commonimportFileIO,Serialization,Typingfromnemo.utilsimportlogging__all__=['NeuralModule','freeze','unfreeze']deffreeze(module:Module)->None:"""Freeze all parameters of ``module`` and snapshot their prior ``requires_grad`` state. The snapshot is stored on ``module._frozen_grad_map`` so a later call to ``unfreeze(..., partial=True)`` can restore the pre-freeze state instead of unconditionally enabling gradients. """grad_map={pname:param.requires_gradforpname,paraminmodule.named_parameters()}forparaminmodule.parameters():param.requires_grad=Falseifnothasattr(module,'_frozen_grad_map'):module._frozen_grad_map=grad_mapelse:module._frozen_grad_map.update(grad_map)module.eval()defunfreeze(module:Module,partial:bool=False)->None:"""Unfreeze parameters of ``module``. If ``partial=True``, restore each parameter's ``requires_grad`` from the snapshot recorded by ``freeze(module)``; otherwise enable gradients on every parameter. The snapshot is cleared in both cases and ``module.train()`` is called. """ifpartialandnothasattr(module,'_frozen_grad_map'):raiseValueError("Cannot unfreeze partially without first freezing the module with `freeze()`")forpname,paraminmodule.named_parameters():ifnotpartial:param.requires_grad=Trueelifpnameinmodule._frozen_grad_map:param.requires_grad=module._frozen_grad_map[pname]else:logging.warning(f"Parameter {pname} not found in list of previously frozen parameters. Unfreezing this parameter.")param.requires_grad=Trueifhasattr(module,'_frozen_grad_map'):delattr(module,'_frozen_grad_map')module.train()
classNeuralModule(Module,Typing,Serialization,FileIO):""" Abstract class offering interface shared between all PyTorch Neural Modules. """@propertydefnum_weights(self):""" Utility property that returns the total number of parameters of NeuralModule. """returnself._num_weights()@torch.jit.ignoredef_num_weights(self):num:int=0forpinself.parameters():ifp.requires_grad:num+=p.numel()returnnum
definput_example(self,max_batch=None,max_dim=None):""" Override this method if random inputs won't work Returns: A tuple sample of valid input data. """returnNone
deffreeze(self)->None:r"""Freeze all params for inference. See :func:`freeze` for details."""freeze(self)
defunfreeze(self,partial:bool=False)->None:"""Unfreeze parameters for training. See :func:`unfreeze` for details. Example: ```python model.encoder.freeze() # caller freezes encoder model.freeze() # freezes everything; encoder snapshot preserved model.unfreeze(partial=True) # decoder unfrozen, encoder stays frozen ``` """unfreeze(self,partial=partial)
@contextmanagerdefas_frozen(self):""" Context manager which temporarily freezes a module, yields control and finally unfreezes the module partially to return to original state. Allows for either total unfreeze or partial unfreeze (if the module was explicitly frozen previously with `freeze()`). The `partial` argument is used to determine whether to unfreeze all parameters or only the parameters that were previously unfrozen prior `freeze()`. Example: with model.as_frozen(): # by default, partial = True # Do something with the model pass # Model's parameters are now back to original state of requires_grad """training_mode=self.trainingself.freeze()try:yieldfinally:self.unfreeze(partial=True)iftraining_mode:self.train()else:self.eval()