From e27144ebfc5c08bb10a40d6c4361225fb33b3bd4 Mon Sep 17 00:00:00 2001 From: Anton Schwaighofer Date: Wed, 9 Feb 2022 15:05:45 +0000 Subject: [PATCH] using plain dataloaders --- InnerEye/ML/SSL/lightning_containers/ssl_container.py | 11 +++++------ .../ML/SSL/lightning_modules/ssl_online_evaluator.py | 2 +- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/InnerEye/ML/SSL/lightning_containers/ssl_container.py b/InnerEye/ML/SSL/lightning_containers/ssl_container.py index 6e714c8c..0bf82654 100644 --- a/InnerEye/ML/SSL/lightning_containers/ssl_container.py +++ b/InnerEye/ML/SSL/lightning_containers/ssl_container.py @@ -151,7 +151,7 @@ class SSLContainer(LightningContainer): model: LightningModule = SimCLRInnerEye(encoder_name=self.ssl_encoder.value, dataset_name=self.ssl_training_dataset_name.value, use_7x7_first_conv_in_resnet=use_7x7_first_conv_in_resnet, - num_samples=self.data_module.num_train_samples, + num_samples=len(self.data_module.dataset_train), batch_size=self.data_module.batch_size, gpus=self.num_gpus_per_node(), num_nodes=self.num_nodes, @@ -172,8 +172,8 @@ class SSLContainer(LightningContainer): f"{SSLTrainingType.BYOL.value}. " f"Found {self.ssl_training_type.value}") model.hparams.update({'ssl_type': self.ssl_training_type.value, - "num_classes": self.data_module.num_classes}) - self.encoder_output_dim = get_encoder_output_dim(model, self.data_module) + "num_classes": 0}) + self.encoder_output_dim = -1 # get_encoder_output_dim(model, self.data_module) return model def get_data_module(self) -> InnerEyeDataModuleTypes: @@ -184,9 +184,8 @@ class SSLContainer(LightningContainer): if hasattr(self, "data_module"): return self.data_module encoder_data_module = self._create_ssl_data_modules(is_ssl_encoder_module=True) - linear_data_module = self._create_ssl_data_modules(is_ssl_encoder_module=False) - return CombinedDataModule(encoder_data_module, linear_data_module, - self.use_balanced_binary_loss_for_linear_head) + # linear_data_module = self._create_ssl_data_modules(is_ssl_encoder_module=False) + return encoder_data_module def _create_ssl_data_modules(self, is_ssl_encoder_module: bool) -> InnerEyeVisionDataModule: """ diff --git a/InnerEye/ML/SSL/lightning_modules/ssl_online_evaluator.py b/InnerEye/ML/SSL/lightning_modules/ssl_online_evaluator.py index 2aa77096..ecaf2fe9 100644 --- a/InnerEye/ML/SSL/lightning_modules/ssl_online_evaluator.py +++ b/InnerEye/ML/SSL/lightning_modules/ssl_online_evaluator.py @@ -124,7 +124,7 @@ class SSLOnlineEvaluatorInnerEye(SSLOnlineEvaluator): Moves batch to device. :param device: device to move the batch to. """ - _, x, y = batch + x, y = batch return x.to(device), y.to(device) def on_train_epoch_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: