Cannot import BaaL integrations

[betairylia]

🐛 Bug

I could not import flash.image.classification.integrations.baal.ActiveLearningDataModule as being described in https://devblog.pytorchlightning.ai/active-learning-made-simple-using-flash-and-baal-2216df6f872c

To Reproduce

Code sample

from flash.image.classification.integrations.baal import ActiveLearningDataModule

gives

ImportError                               Traceback (most recent call last)
Input In [3], in <cell line: 1>()
----> 1 from flash.image.classification.integrations.baal import ActiveLearningDataModule

File ~/anaconda3/envs/AL/lib/python3.9/site-packages/flash/image/classification/integrations/baal/__init__.py:2, in <module>
      1 from flash.image.classification.integrations.baal.data import ActiveLearningDataModule  # noqa F401
----> 2 from flash.image.classification.integrations.baal.loop import ActiveLearningLoop

File ~/anaconda3/envs/AL/lib/python3.9/site-packages/flash/image/classification/integrations/baal/loop.py:28, in <module>
     26 from flash.core.utilities.stages import RunningStage
     27 from flash.image.classification.integrations.baal.data import ActiveLearningDataModule
---> 28 from flash.image.classification.integrations.baal.dropout import InferenceMCDropoutTask
     30 if _PL_GREATER_EQUAL_1_4_0:
     31     from pytorch_lightning.loops import Loop

File ~/anaconda3/envs/AL/lib/python3.9/site-packages/flash/image/classification/integrations/baal/dropout.py:21, in <module>
     18 from flash.core.utilities.imports import _BAAL_AVAILABLE
     20 if _BAAL_AVAILABLE:
---> 21     from baal.bayesian.dropout import _patch_dropout_layers
     24 class InferenceMCDropoutTask(flash.Task):
     25     def __init__(self, module: flash.Task, inference_iteration: int):

ImportError: cannot import name '_patch_dropout_layers' from 'baal.bayesian.dropout' (.../baal/bayesian/dropout.py)

Expected behavior

Modules got imported without errors.

Environment

• OS (e.g., Linux): Ubuntu 18.04
• Python version: 3.9
• PyTorch/Lightning/Flash Version (e.g., 1.10/1.5/0.7): 1.11/1.6.1/0.7.3
• GPU models and configuration:
• Any other relevant information: BaaL 1.5.2

Additional context

Seems baal team has removed _patch_dropout_layers: baal-org/baal@1006861 (baal/bayesian/dropout.py).
Is this relevant or I'm doing it in a wrong way?

[krshrimali]

Hi, @betairylia - Thanks for creating this issue and also for the additional context. Apart from the _patch_dropout_layers, there are a few other things that have changed in baal. Our CI currently tests against 1.3.2 version of baal, and I'll be working on making Flash work on baal's latest release.

Meanwhile, if it's urgent, and you need to make it work, could you downgrade the baal version to 1.3.2? pip install baal==1.3.2.

[betairylia]

Hi @krshrimali, Thank you for your reply and the details.
Have a good day!