Lightning Trainer/PyTorch Lightning 1.7.0 support + CI Fixes (JIT Tracing and Functions to Classes conversion)

.azure-pipelines/testing-template.yml

@@ -36,7 +36,8 @@ jobs:
 
       - bash: |
           # python -m pip install "pip==20.1"
-          if [ "${{config}}" == "icevision" ]; then pip install '.[image]' icevision effdet icedata; elif [ "${{config}}" == "vissl" ]; then pip install '.[image]'; else pip install '.[${{config}}]'; fi
+          if [ "${{config}}" == "icevision" ]; then pip install '.[image]' icevision icedata; elif [ "${{config}}" == "vissl" ]; then pip install '.[image]'; else pip install '.[${{config}}]'; fi
+          if [ "${{config}}" == "icevision" ]; then pip install git+https://github.com/rwightman/efficientdet-pytorch.git@79d26d8982b9f8e1f27d9f896e38012ac250fd26; fi
           pip install '.[test]' --upgrade-strategy only-if-needed
           pip list
         displayName: 'Install dependencies'

.github/workflows/ci-testing.yml

@@ -33,6 +33,7 @@ jobs:
           - { os: 'ubuntu-20.04', python-version: 3.9, requires: 'latest', release: 'pre', topic: [ 'core' ] }
           - { os: 'ubuntu-20.04', python-version: 3.9, requires: 'latest', release: 'stable', topic: [ 'image' ] }
           - { os: 'ubuntu-20.04', python-version: 3.9, requires: 'latest', release: 'stable', topic: [ 'image','image_extras' ] }
+          - { os: 'ubuntu-20.04', python-version: 3.9, requires: 'latest', release: 'stable', topic: [ 'image','image_extras_effdet' ] }
           - { os: 'ubuntu-20.04', python-version: 3.9, requires: 'latest', release: 'stable', topic: [ 'image','image_extras_baal' ] }
           - { os: 'ubuntu-20.04', python-version: 3.9, requires: 'latest', release: 'stable', topic: [ 'video' ] }
           - { os: 'ubuntu-20.04', python-version: 3.9, requires: 'latest', release: 'stable', topic: [ 'video','video_extras' ] }
@@ -128,7 +129,8 @@ jobs:
       if: contains( matrix.topic , 'serve' )
       run: |
         sudo apt-get install libsndfile1
-        pip install '.[all,audio]' icevision effdet --upgrade
+        pip install '.[all,audio]' icevision --upgrade
+        pip install -e git+https://github.com/rwightman/efficientdet-pytorch.git@79d26d8982b9f8e1f27d9f896e38012ac250fd26#egg=effdet
 
     - name: Install audio test dependencies
       if: contains( matrix.topic , 'audio' )

flash/core/integrations/icevision/adapter.py

@@ -110,6 +110,27 @@ def _wrap_collate_fn(collate_fn, samples):
             DataKeys.METADATA: metadata,
         }
 
+    def change_collate_fn_dataloader(self, data_loader):
+        # Starting PL 1.7.0 - changing attributes after the DataLoader is initialized - will not work
+        # So we manually set the collate_fn to the wrapped version, for now.
+        new_kwargs = getattr(data_loader, "__pl_saved_kwargs", None)
+        new_collate_fn = functools.partial(self._wrap_collate_fn, data_loader.collate_fn)
+        if new_kwargs:
+            new_kwargs["collate_fn"] = new_collate_fn
+            setattr(data_loader, "__pl_saved_kwargs", new_kwargs)
+        data_loader.collate_fn = new_collate_fn
+        return data_loader
+
+    def update_collate_fn_dataloader(self, new_collate_fn, data_loader):
+        # Starting PL 1.7.0 - changing attributes after the DataLoader is initialized - will not work
+        # So we manually update the collate_fn for the dataloader, for now.
+        new_kwargs = getattr(data_loader, "__pl_saved_kwargs", None)
+        if new_kwargs:
+            new_kwargs["collate_fn"] = new_collate_fn
+            setattr(data_loader, "__pl_saved_kwargs", new_kwargs)
+        data_loader.collate_fn = new_collate_fn
+        return data_loader
+
     def process_train_dataset(
         self,
         dataset: InputBase,
@@ -134,12 +155,13 @@ def process_train_dataset(
             persistent_workers=persistent_workers,
         )
 
-        data_loader.collate_fn = functools.partial(self._wrap_collate_fn, data_loader.collate_fn)
+        data_loader = self.change_collate_fn_dataloader(data_loader)
 
         input_transform = input_transform or self.input_transform
         if input_transform is not None:
             input_transform.inject_collate_fn(data_loader.collate_fn)
-            data_loader.collate_fn = create_worker_input_transform_processor(RunningStage.TRAINING, input_transform)
+            new_collate_fn = create_worker_input_transform_processor(RunningStage.TRAINING, input_transform)
+            data_loader = self.update_collate_fn_dataloader(new_collate_fn, data_loader)
         return data_loader
 
     def process_val_dataset(
@@ -166,12 +188,13 @@ def process_val_dataset(
             persistent_workers=persistent_workers,
         )
 
-        data_loader.collate_fn = functools.partial(self._wrap_collate_fn, data_loader.collate_fn)
+        data_loader = self.change_collate_fn_dataloader(data_loader)
 
         input_transform = input_transform or self.input_transform
         if input_transform is not None:
             input_transform.inject_collate_fn(data_loader.collate_fn)
-            data_loader.collate_fn = create_worker_input_transform_processor(RunningStage.VALIDATING, input_transform)
+            new_collate_fn = create_worker_input_transform_processor(RunningStage.VALIDATING, input_transform)
+            data_loader = self.update_collate_fn_dataloader(new_collate_fn, data_loader)
         return data_loader
 
     def process_test_dataset(
@@ -198,12 +221,13 @@ def process_test_dataset(
             persistent_workers=persistent_workers,
         )
 
-        data_loader.collate_fn = functools.partial(self._wrap_collate_fn, data_loader.collate_fn)
+        data_loader = self.change_collate_fn_dataloader(data_loader)
 
         input_transform = input_transform or self.input_transform
         if input_transform is not None:
             input_transform.inject_collate_fn(data_loader.collate_fn)
-            data_loader.collate_fn = create_worker_input_transform_processor(RunningStage.TESTING, input_transform)
+            new_collate_fn = create_worker_input_transform_processor(RunningStage.TESTING, input_transform)
+            data_loader = self.update_collate_fn_dataloader(new_collate_fn, data_loader)
         return data_loader
 
     def process_predict_dataset(
@@ -230,12 +254,13 @@ def process_predict_dataset(
             persistent_workers=persistent_workers,
         )
 
-        data_loader.collate_fn = functools.partial(self._wrap_collate_fn, data_loader.collate_fn)
+        data_loader = self.change_collate_fn_dataloader(data_loader)
 
         input_transform = input_transform or self.input_transform
         if input_transform is not None:
             input_transform.inject_collate_fn(data_loader.collate_fn)
-            data_loader.collate_fn = create_worker_input_transform_processor(RunningStage.PREDICTING, input_transform)
+            new_collate_fn = create_worker_input_transform_processor(RunningStage.PREDICTING, input_transform)
+            data_loader = self.update_collate_fn_dataloader(new_collate_fn, data_loader)
         return data_loader
 
     def training_step(self, batch, batch_idx) -> Any:

flash/core/utilities/imports.py

@@ -301,7 +301,9 @@ def _import_module(self):
     topic = os.environ["FLASH_TEST_TOPIC"]
     _CORE_TESTING = topic == "core"
     _IMAGE_TESTING = topic == "image"
-    _IMAGE_EXTRAS_TESTING = topic == "image,image_extras" or topic == "icevision" or topic == "vissl"
+    _IMAGE_EXTRAS_TESTING = (
+        topic == "image,image_extras,image_extras_effdet" or topic == "icevision" or topic == "vissl"
+    )
     _VIDEO_TESTING = topic == "video"
     _VIDEO_EXTRAS_TESTING = topic == "video,video_extras"
     _TABULAR_TESTING = topic == "tabular"

flash/core/utilities/lightning_cli.py

@@ -97,7 +97,8 @@ def add_lightning_class_args(
             lightning_class = class_from_function(lightning_class)
 
         if inspect.isclass(lightning_class) and issubclass(
-            cast(type, lightning_class), (Trainer, LightningModule, LightningDataModule, Callback)
+            cast(type, lightning_class),
+            (Trainer, LightningModule, LightningDataModule, Callback, ClassFromFunctionBase),
         ):
             if issubclass(cast(type, lightning_class), Callback):
                 self.callback_keys.append(nested_key)

flash/image/classification/integrations/baal/data.py

@@ -180,7 +180,7 @@ def label(self, probabilities: List[Tensor] = None, indices=None):
             raise MisconfigurationException(
                 "The `probabilities` and `indices` are mutually exclusive, pass only of one them."
             )
-        if probabilities is not None:
+        if probabilities is not None and len(probabilities) != 0:
             probabilities = torch.cat([p[0].unsqueeze(0) for p in probabilities], dim=0)
             uncertainties = self.heuristic.get_uncertainties(probabilities)
             indices = np.argsort(uncertainties)

flash/image/classification/integrations/baal/loop.py

@@ -91,7 +91,7 @@ def on_run_start(self, *args: Any, **kwargs: Any) -> None:
         assert isinstance(self.trainer.datamodule, ActiveLearningDataModule)
         if self._datamodule_state_dict is not None:
             self.trainer.datamodule.load_state_dict(self._datamodule_state_dict)
-        self.trainer.predict_loop._return_predictions = True
+        self.trainer.predict_loop.return_predictions = True
         self._lightning_module = self.trainer.lightning_module
         self._model_state_dict = deepcopy(self._lightning_module.state_dict())
         self.inference_model = InferenceMCDropoutTask(self._lightning_module, self.inference_iteration)
@@ -165,21 +165,18 @@ def _connect(self, model: LightningModule):
     def _reset_fitting(self):
         self.trainer.state.fn = TrainerFn.FITTING
         self.trainer.training = True
-        self.trainer.lightning_module.on_train_dataloader()
         self._connect(self._lightning_module)
         self.fit_loop.epoch_progress = Progress()
 
     def _reset_predicting(self):
         self.trainer.state.fn = TrainerFn.PREDICTING
         self.trainer.predicting = True
-        self.trainer.lightning_module.on_predict_dataloader()
         self._connect(self.inference_model)
 
     def _reset_testing(self):
         self.trainer.state.fn = TrainerFn.TESTING
         self.trainer.state.status = TrainerStatus.RUNNING
         self.trainer.testing = True
-        self.trainer.lightning_module.on_test_dataloader()
         self._connect(self._lightning_module)
 
     def _reset_dataloader_for_stage(self, running_state: RunningStage):

flash/image/embedding/model.py

@@ -154,8 +154,8 @@ def on_train_start(self) -> None:
     def on_train_epoch_end(self) -> None:
         self.adapter.on_train_epoch_end()
 
-    def on_train_batch_end(self, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
-        self.adapter.on_train_batch_end(outputs, batch, batch_idx, dataloader_idx)
+    def on_train_batch_end(self, outputs: Any, batch: Any, batch_idx: int, *args) -> None:
+        self.adapter.on_train_batch_end(outputs, batch, batch_idx, *args)
 
     @classmethod
     @requires("image", "vissl", "fairscale")

flash/image/embedding/vissl/hooks.py

@@ -96,7 +96,7 @@ def on_train_start(self) -> None:
         for hook in self.hooks:
             hook.on_start(self.task)
 
-    def on_train_batch_end(self, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
+    def on_train_batch_end(self, outputs: Any, batch: Any, batch_idx: int, *args) -> None:
         self.task.iteration += 1
 
     def on_train_epoch_end(self) -> None:

requirements/datatype_image_extras.txt

@@ -4,7 +4,6 @@ classy_vision
 vissl>=0.1.5
 icevision>=0.8
 icedata
-effdet
 albumentations
 learn2learn
 fastface

tests/helpers/task_tester.py

@@ -102,8 +102,10 @@ def _test_jit_trace(self, tmpdir):
     path = os.path.join(tmpdir, "test.pt")
 
     model = self.instantiated_task
+    trainer = self.instantiated_trainer
     model.eval()
 
+    model.trainer = trainer
     model = torch.jit.trace(model, self.example_forward_input)
 
     torch.jit.save(model, path)
@@ -117,8 +119,10 @@ def _test_jit_script(self, tmpdir):
     path = os.path.join(tmpdir, "test.pt")
 
     model = self.instantiated_task
+    trainer = self.instantiated_trainer
     model.eval()
 
+    model.trainer = trainer
     model = torch.jit.script(model)
 
     torch.jit.save(model, path)
@@ -261,10 +265,17 @@ class TaskTester(metaclass=TaskTesterMeta):
         "test_cli": [pytest.mark.parametrize("extra_args", [[]])],
     }
 
+    trainer_args: Tuple = ()
+    trainer_kwargs: Dict = {}
+
     @property
     def instantiated_task(self):
         return self.task(*self.task_args, **self.task_kwargs)
 
+    @property
+    def instantiated_trainer(self):
+        return flash.Trainer(*self.trainer_args, **self.trainer_kwargs)
+
     @property
     def example_forward_input(self):
         pass