Fix invalid value for weights_summary

pytorch_lightning/core/lightning.py

@@ -14,15 +14,15 @@
 
 """nn.Module with additional great features."""
 
-from abc import ABC
-from argparse import Namespace
 import collections
 import copy
 import inspect
 import os
-from pathlib import Path
 import re
 import tempfile
+from abc import ABC
+from argparse import Namespace
+from pathlib import Path
 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
 import torch
@@ -1327,9 +1327,17 @@ def tbptt_split_batch(self, batch, split_size):
 
  return splits
 
- def summarize(self, mode: str = ModelSummary.MODE_DEFAULT) -> ModelSummary:
- model_summary = ModelSummary(self, mode=mode)
- log.info("\n" + str(model_summary))
+ def summarize(self, mode: Optional[str] = ModelSummary.MODE_DEFAULT) -> Optional[ModelSummary]:
+ model_summary = None
+
+ if mode in ModelSummary.MODES:
+ model_summary = ModelSummary(self, mode=mode)
+ log.info("\n" + str(model_summary))
+ elif mode is not None:
+ raise MisconfigurationException(
+ f"`mode` can be None, {', '.join(ModelSummary.MODES)}, got {mode}"
+ )
+
  return model_summary
 
  def freeze(self) -> None:

pytorch_lightning/trainer/trainer.py

@@ -311,7 +311,6 @@ def __init__(
  self.plugin_connector = PluginConnector(self)
 
  # training state
- self.weights_summary = weights_summary
  self.model = None
  self.shown_warnings = set()
 
@@ -374,7 +373,8 @@ def __init__(
  max_steps,
  min_steps,
  num_sanity_val_steps,
- automatic_optimization
+ automatic_optimization,
+ weights_summary,
  )
  self.evaluation_loop.on_trainer_init()
 

pytorch_lightning/trainer/training_loop.py

@@ -49,7 +49,14 @@ def __init__(self, trainer):
  self._cur_grad_norm_dict = None
 
  def on_trainer_init(
- self, max_epochs, min_epochs, max_steps, min_steps, num_sanity_val_steps, automatic_optimization
+ self,
+ max_epochs,
+ min_epochs,
+ max_steps,
+ min_steps,
+ num_sanity_val_steps,
+ automatic_optimization,
+ weights_summary,
  ):
  self.trainer.global_step = 0
  self.trainer.current_epoch = 0
@@ -73,6 +80,12 @@ def on_trainer_init(
  else:
  self.trainer.num_sanity_val_steps = num_sanity_val_steps
 
+ self.trainer.weights_summary = weights_summary
+ if weights_summary is not None and weights_summary not in ModelSummary.MODES:
+ raise MisconfigurationException(
+ f"`weights_summary` can be None, {', '.join(ModelSummary.MODES)}, got {weights_summary}"
+ )
+
  @property
  def num_optimizers(self):
  num_optimizers = len(self.get_optimizers_iterable())
@@ -161,11 +174,8 @@ def setup_training(self, model: LightningModule):
  ref_model.on_pretrain_routine_start()
 
  # print model summary
- if self.trainer.is_global_zero and self.trainer.weights_summary is not None and not self.trainer.testing:
- if self.trainer.weights_summary in ModelSummary.MODES:
- ref_model.summarize(mode=self.trainer.weights_summary)
- else:
- raise MisconfigurationException("weights_summary can be None, " + ", ".join(ModelSummary.MODES))
+ if self.trainer.is_global_zero and not self.trainer.testing:
+ ref_model.summarize(mode=self.trainer.weights_summary)
 
  # track model now.
  # if cluster resets state, the model will update with the saved weights

tests/core/test_memory.py

@@ -15,8 +15,9 @@
 import torch
 import torch.nn as nn
 
-from pytorch_lightning import LightningModule
+from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.core.memory import UNKNOWN_SIZE, ModelSummary
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.base.models import ParityModuleRNN
 
 
@@ -68,6 +69,15 @@ def forward(self, x):
  return self.reduce(self.embed(x))
 
 
+def test_invalid_weights_summmary():
+ """ Test that invalid value for weights_summary raises an error. """
+ with pytest.raises(MisconfigurationException, match='`mode` can be None, .* got temp'):
+ UnorderedModel().summarize(mode='temp')
+
+ with pytest.raises(MisconfigurationException, match='`weights_summary` can be None, .* got temp'):
+ Trainer(weights_summary='temp')
+
+
 @pytest.mark.parametrize(['mode'], [
  pytest.param(ModelSummary.MODE_FULL),
  pytest.param(ModelSummary.MODE_TOP),