Decouple DataModules from Models - SemSegment

pl_bolts/models/vision/segmentation.py

@@ -9,13 +9,12 @@
 
 class SemSegment(pl.LightningModule):
  def __init__(
- self,
- datamodule: pl.LightningDataModule = None,
- lr: float = 0.01,
- num_classes: int = 19,
- num_layers: int = 5,
- features_start: int = 64,
- bilinear: bool = False
+ self,
+ lr: float = 0.01,
+ num_classes: int = 19,
+ num_layers: int = 5,
+ features_start: int = 64,
+ bilinear: bool = False
  ):
  """
  Basic model for semantic segmentation. Uses UNet architecture by default.
@@ -29,17 +28,13 @@ def __init__(
  - `Annika Brundyn <https://github.com/annikabrundyn>`_
 
  Args:
- datamodule: LightningDataModule
  num_layers: number of layers in each side of U-net (default 5)
  features_start: number of features in first layer (default 64)
  bilinear: whether to use bilinear interpolation (True) or transposed convolutions (default) for upsampling.
  lr: learning (default 0.01)
  """
  super().__init__()
 
- assert datamodule
- self.datamodule = datamodule
-
  self.num_classes = num_classes
  self.num_layers = num_layers
  self.features_start = features_start
@@ -84,7 +79,6 @@ def configure_optimizers(self):
  @staticmethod
  def add_model_specific_args(parent_parser):
  parser = ArgumentParser(parents=[parent_parser], add_help=False)
- parser.add_argument("--batch_size", type=int, default=16, help="size of the batches")
  parser.add_argument("--lr", type=float, default=0.01, help="adam: learning rate")
  parser.add_argument("--num_layers", type=int, default=5, help="number of layers on u-net")
  parser.add_argument("--features_start", type=float, default=64, help="number of features in first layer")
@@ -100,23 +94,24 @@ def cli_main():
  pl.seed_everything(1234)
 
  parser = ArgumentParser()
-
  # trainer args
  parser = pl.Trainer.add_argparse_args(parser)
-
  # model args
  parser = SemSegment.add_model_specific_args(parser)
+ # datamodule args
+ parser = KittiDataModule.add_argparse_args(parser)
+
  args = parser.parse_args()
 
  # data
  dm = KittiDataModule(args.data_dir).from_argparse_args(args)
 
  # model
- model = SemSegment(**args.__dict__, datamodule=dm)
+ model = SemSegment(**args.__dict__)
 
  # train
  trainer = pl.Trainer().from_argparse_args(args)
- trainer.fit(model)
+ trainer.fit(model, dm)
 
 
 if __name__ == '__main__':

tests/models/test_vision.py

@@ -66,10 +66,10 @@ def train_dataloader(self):
 
  dm = DummyDataModule()
 
- model = SemSegment(datamodule=dm, num_classes=19)
+ model = SemSegment(num_classes=19)
 
  trainer = pl.Trainer(fast_dev_run=True, max_epochs=1)
- trainer.fit(model)
+ trainer.fit(model, dm)
  loss = trainer.progress_bar_dict['loss']
 
  assert float(loss) > 0