Added model_transforms in models

src/ptbench/models/alexnet.py

@@ -73,6 +73,12 @@ class Alexnet(pl.LightningModule):
 
         self.name = "alexnet"
 
+        self.model_transforms = [
+            torchvision.transforms.ToPILImage(),
+            torchvision.transforms.Lambda(lambda x: x.convert("RGB")),
+            torchvision.transforms.ToTensor(),
+        ]
+
         self._train_loss = train_loss
         self._validation_loss = (
             validation_loss if validation_loss is not None else train_loss

src/ptbench/models/densenet.py

@@ -71,6 +71,12 @@ class Densenet(pl.LightningModule):
 
         self.name = "densenet-121"
 
+        self.model_transforms = [
+            torchvision.transforms.ToPILImage(),
+            torchvision.transforms.Lambda(lambda x: x.convert("RGB")),
+            torchvision.transforms.ToTensor(),
+        ]
+
         self._train_loss = train_loss
         self._validation_loss = (
             validation_loss if validation_loss is not None else train_loss

src/ptbench/models/pasa.py

@@ -72,6 +72,8 @@ class Pasa(pl.LightningModule):
 
         self.name = "pasa"
 
+        self.model_transforms = []
+
         self._train_loss = train_loss
         self._validation_loss = (
             validation_loss if validation_loss is not None else train_loss

src/ptbench/scripts/predict.py

@@ -100,6 +100,8 @@ def predict(
     from ..utils.plot import relevance_analysis_plot
 
     datamodule.set_chunk_size(batch_size, 1)
+    datamodule.model_transforms = model.model_transforms
+
 
     logger.info(f"Loading checkpoint from {weight}")
     model = model.load_from_checkpoint(weight, strict=False)

src/ptbench/scripts/train.py

@@ -248,6 +248,7 @@ def train(
     datamodule.drop_incomplete_batch = drop_incomplete_batch
     datamodule.cache_samples = cache_samples
     datamodule.parallel = parallel
+    datamodule.model_transforms = model.model_transforms
 
     datamodule.prepare_data()
     datamodule.setup(stage="fit")