Update lightning_template model to save hparams. (Lightning-AI#2665)

Author: romesco

pl_examples/models/lightning_template.py

@@ -24,51 +24,42 @@ class LightningTemplateModel(LightningModule):
 
         >>> # define simple Net for MNIST dataset
         >>> params = dict(
-        ...     drop_prob=0.2,
-        ...     batch_size=2,
         ...     in_features=28 * 28,
+        ...     hidden_dim=1000,
+        ...     out_features=10,
+        ...     drop_prob=0.2,
         ...     learning_rate=0.001 * 8,
-        ...     optimizer_name='adam',
+        ...     batch_size=2,
         ...     data_root='./datasets',
-        ...     out_features=10,
         ...     num_workers=4,
-        ...     hidden_dim=1000,
         ... )
         >>> model = LightningTemplateModel(**params)
     """
 
     def __init__(self,
-                 drop_prob: float = 0.2,
-                 batch_size: int = 2,
                  in_features: int = 28 * 28,
+                 hidden_dim: int = 1000,
+                 out_features: int = 10,
+                 drop_prob: float = 0.2,
                  learning_rate: float = 0.001 * 8,
-                 optimizer_name: str = 'adam',
+                 batch_size: int = 2,
                  data_root: str = './datasets',
-                 out_features: int = 10,
                  num_workers: int = 4,
-                 hidden_dim: int = 1000,
                  **kwargs
                  ):
         # init superclass
         super().__init__()
+        # save all variables in __init__ signature to self.hparams
+        self.save_hyperparameters()
 
-        self.num_workers = num_workers
-        self.drop_prob = drop_prob
-        self.batch_size = batch_size
-        self.in_features = in_features
-        self.learning_rate = learning_rate
-        self.optimizer_name = optimizer_name
-        self.data_root = data_root
-        self.out_features = out_features
-        self.hidden_dim = hidden_dim
+        self.c_d1 = nn.Linear(in_features=self.hparams.in_features,
+                              out_features=self.hparams.hidden_dim)
+        self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
+        self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
 
-        self.c_d1 = nn.Linear(in_features=self.in_features,
-                              out_features=self.hidden_dim)
-        self.c_d1_bn = nn.BatchNorm1d(self.hidden_dim)
-        self.c_d1_drop = nn.Dropout(self.drop_prob)
+        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim,
 
-        self.c_d2 = nn.Linear(in_features=self.hidden_dim,
-                              out_features=self.out_features)
+                              out_features=self.hparams.out_features)
 
         self.example_input_array = torch.zeros(2, 1, 28, 28)
 
@@ -140,27 +131,27 @@ def configure_optimizers(self):
         Return whatever optimizers and learning rate schedulers you want here.
         At least one optimizer is required.
         """
-        optimizer = optim.Adam(self.parameters(), lr=self.learning_rate)
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
         scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
         return [optimizer], [scheduler]
 
     def prepare_data(self):
-        MNIST(self.data_root, train=True, download=True, transform=transforms.ToTensor())
-        MNIST(self.data_root, train=False, download=True, transform=transforms.ToTensor())
+        MNIST(self.hparams.data_root, train=True, download=True, transform=transforms.ToTensor())
+        MNIST(self.hparams.data_root, train=False, download=True, transform=transforms.ToTensor())
 
     def setup(self, stage):
         transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
-        self.mnist_train = MNIST(self.data_root, train=True, download=False, transform=transform)
-        self.mnist_test = MNIST(self.data_root, train=False, download=False, transform=transform)
+        self.mnist_train = MNIST(self.hparams.data_root, train=True, download=False, transform=transform)
+        self.mnist_test = MNIST(self.hparams.data_root, train=False, download=False, transform=transform)
 
     def train_dataloader(self):
-        return DataLoader(self.mnist_train, batch_size=self.batch_size, num_workers=self.num_workers)
+        return DataLoader(self.mnist_train, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)
 
     def val_dataloader(self):
-        return DataLoader(self.mnist_test, batch_size=self.batch_size, num_workers=self.num_workers)
+        return DataLoader(self.mnist_test, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)
 
     def test_dataloader(self):
-        return DataLoader(self.mnist_test, batch_size=self.batch_size, num_workers=self.num_workers)
+        return DataLoader(self.mnist_test, batch_size=self.hparams.batch_size, num_workers=self.hparams.num_workers)
 
     @staticmethod
     def add_model_specific_args(parent_parser, root_dir):  # pragma: no-cover
@@ -174,18 +165,17 @@ def add_model_specific_args(parent_parser, root_dir):  # pragma: no-cover
 
         # network params
         parser.add_argument('--in_features', default=28 * 28, type=int)
-        parser.add_argument('--out_features', default=10, type=int)
-        # use 500 for CPU, 50000 for GPU to see speed difference
         parser.add_argument('--hidden_dim', default=50000, type=int)
+        # use 500 for CPU, 50000 for GPU to see speed difference
+        parser.add_argument('--out_features', default=10, type=int)
         parser.add_argument('--drop_prob', default=0.2, type=float)
-        parser.add_argument('--learning_rate', default=0.001, type=float)
-        parser.add_argument('--num_workers', default=4, type=int)
 
         # data
         parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)
+        parser.add_argument('--num_workers', default=4, type=int)
 
         # training params (opt)
         parser.add_argument('--epochs', default=20, type=int)
-        parser.add_argument('--optimizer_name', default='adam', type=str)
         parser.add_argument('--batch_size', default=64, type=int)
+        parser.add_argument('--learning_rate', default=0.001, type=float)
         return parser