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

krshrimali · 2022-08-04T02:56:16Z

What does this PR do?

This PR adds support for PL 1.7.0 to Flash, and fixes a few issues with the CI. Summary:

CI has been failing for Flash CI for the case when a function is converted to ClassFromFunction class. We were not checking if its type is subclass of CustomClassFunctionBase.
For JIT tracing, the model needs to be attached to the Trainer class, else it will give an error similar to: ImageClassifier is not attached to Trainer.
Baal integration fails on recent Lightning Trainer versions. on_predict_dataloader, on_train_dataloader, on_test_dataloader were deprecated, and removed in PL 1.7. I'm not very sure if we need them, removing that works?
For BaaL integration, probabilities need to have a check if it's not empty else torch.cat will fail.
In PL 1.7.0, DataLoader class won't change the attributes except when they are in the __init__ method. Flash changes the collate function after the DataLoader is initialized - this was causing multiple errors. This PR fixes it.

This PR attempts to fix both of the issues. Note that we should document that the best way to JIT trace the model in Flash is to:

# 2. Build the task
model = ImageClassifier(backbone="resnet18", labels=datamodule.labels)

# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count())

# Attach trainer to the model object
model.trainer = trainer
trace = torch.jit.trace(model, torch.randn((3, 196, 196))

Issue: #1413

Before submitting

Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
Did you read the contributor guideline, Pull Request section?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes?
Did you write any new necessary tests? [not needed for typos/docs]
Did you verify new and existing tests pass locally with your changes?
If you made a notable change (that affects users), did you update the CHANGELOG?

PR review

Is this pull request ready for review? (if not, please submit in draft mode)

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

for more information, see https://pre-commit.ci

codecov · 2022-08-04T03:00:25Z

Codecov Report

Merging #1410 (8e03fd7) into master (8737ee8) will decrease coverage by 0.00%.
The diff coverage is 85.00%.

@@            Coverage Diff             @@
##           master    #1410      +/-   ##
==========================================
- Coverage   91.77%   91.76%   -0.01%     
==========================================
  Files         286      286              
  Lines       12857    12861       +4     
==========================================
+ Hits        11799    11802       +3     
- Misses       1058     1059       +1

Flag	Coverage Δ
unittests	`91.76% <85.00%> (-0.01%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
flash/core/utilities/lightning_cli.py	`98.16% <ø> (ø)`
flash/core/integrations/icevision/adapter.py	`94.78% <80.00%> (-1.52%)`	⬇️
...ash/image/classification/integrations/baal/data.py	`87.96% <100.00%> (ø)`
...ash/image/classification/integrations/baal/loop.py	`94.81% <100.00%> (-0.12%)`	⬇️
flash/image/embedding/model.py	`98.48% <100.00%> (ø)`
flash/image/embedding/vissl/hooks.py	`98.24% <100.00%> (ø)`
flash/text/question_answering/model.py	`94.55% <0.00%> (+0.68%)`	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

…AI/lightning-flash into flash-ci/functions-to-classes

for more information, see https://pre-commit.ci

requirements/datatype_text.txt

for more information, see https://pre-commit.ci

carmocca · 2022-08-10T15:18:12Z

tests/helpers/task_tester.py

    model.eval()

+    model.trainer = trainer


See Lightning-AI/pytorch-lightning#14036 (comment) for a workaround

Just curious, anything wrong with the current workaround?

for more information, see https://pre-commit.ci

…AI/lightning-flash into flash-ci/functions-to-classes

krshrimali · 2022-08-12T10:47:35Z

About the 5 failures left in the CI:

CI testing / pytest (ubuntu-20.04, 3.9, latest, stable, graph) (pull_request): This failure is flaky, and has been fixed in PL. Once there is another release of PL, we should be good.
CI testing / pytest (ubuntu-20.04, 3.9, latest, stable, audio) (pull_request): PR Remove sd2 sound extensions from supported audio extensions #1409 is opened to fix it.
Lightning-AI.lightning-flash (Special): Horovod versioning error, but I also see /usr/bin/pip not found error. I plan to fix this in a separate PR, doesn't seem to be very urgent as the first step was to support PL 1.7 with Flash.

Thanks to @otaj for his help with a few failures related to PL 1.7, it was very helpful. ❤️ Also @rohitgr7 for a fix! 🔥

cc: @ethanwharris @Borda @carmocca

ethanwharris

Great work @krshrimali ! Some very painful bugs here. Couple of suggestions. Also, let's add some things to the CHANGELOG here 😃

ethanwharris · 2022-08-12T10:45:01Z

flash/core/integrations/icevision/adapter.py

+    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


Seems like a lot of duplication. Seems like change_collate_fn_dataloader is equivalent to:

self.update_collate_fn_dataloader(functools.partial(self._wrap_collate_fn, data_loader.collate_fn), data_loader)

Maybe better to just do that and remove change_collate_fn_dataloader

@ethanwharris, you are right, this could be merged into only one function. Btw, while technically not a bug on PL side, it is definitely a regression on user experience regarding dataloaders and I'm sorry @krshrimali had to run into that. I suspect though, he will not be the last one, so I will send a PR today/early next week to PL to improve UX in there

@ethanwharris - Agreed. Fixed!

@otaj - Thank you for your help with this, please don't be sorry - looking forward to your PR which improves the UX further. Thank you again!

flash/core/integrations/icevision/adapter.py

flash/core/utilities/imports.py

ethanwharris · 2022-08-12T10:50:05Z

.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' ] }


This seems to have no effect as it will lead to just image deps being installed and effdet is installed only for serving. I think it would be cleaner to not add a new job but just set the effdet version in requirements/datatype_image_extras to be the one from github since we don't expect users to install the image_extras extras anyway.

I'm sorry, should have mentioned - I tried putting the version of effdet in datatype_image_extras.txt and doing pip install -r requirements/datatype_image_extras.txt fails miserably :/ Lots of version conflicts. Hence I thought to create a new file, and put effdet there.

From the offline conversation with @ethanwharris - we decided to limit PyTorch version to <1.12 until icevision supports effdet 0.3.0.

for more information, see https://pre-commit.ci

.github/workflows/ci-testing.yml

flash/core/integrations/icevision/adapter.py

for more information, see https://pre-commit.ci

krshrimali · 2022-08-12T13:08:42Z

Hi, @ethanwharris - The 5 failures are expected in the CI, only 1 of these is something that I haven't investigated yet. Fixes for others are either already merged in PL, or we have a PR opened in Flash. Could you please take a look again, and help merge if it all looks good? Thanks!

.azure-pipelines/testing-template.yml

ethanwharris

LGTM 🚀

.azure-pipelines/testing-template.yml

ethanwharris · 2022-08-12T13:13:57Z

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


Not for this PR, but we should figure out what to do with the baal loop. It is so bound to a particular PL and Baal version that I'm not sure it makes sense to have it as part of a framework. Maybe it would be better as a tutorial or in bolts? cc @otaj

…AI/lightning-flash into flash-ci/functions-to-classes

Fix Flash CI: in case of converting functions to classes

d22295e

krshrimali requested review from ethanwharris, Borda, tchaton, justusschock and kaushikb11 as code owners August 4, 2022 02:56

[pre-commit.ci] auto fixes from pre-commit.com hooks

b725620

for more information, see https://pre-commit.ci

krshrimali changed the title ~~Fix Flash CI: in case of converting functions to classes~~ Fix Flash CLI: in case of converting functions to classes Aug 4, 2022

krshrimali marked this pull request as draft August 4, 2022 03:03

krshrimali added 2 commits August 4, 2022 09:31

Attach Trainer to model for jit tracing

6d2c636

Merge branch 'flash-ci/functions-to-classes' of github.com:Lightning-…

f25eb1f

…AI/lightning-flash into flash-ci/functions-to-classes

krshrimali changed the title ~~Fix Flash CLI: in case of converting functions to classes~~ Fix Flash CLI: in case of converting functions to classes + JIT tracing tests Aug 4, 2022

krshrimali added 2 commits August 4, 2022 11:05

on_<>_dataloader were deprecated, and were removed in v1.7 - remove

41fd845

Fix jit_script test

df9f48b

krshrimali mentioned this pull request Aug 5, 2022

CI: Update docker image #1411

Merged

4 tasks

rohitgr7 and others added 3 commits August 8, 2022 14:54

fix

4c8c0fb

Try upgrading torchtext to 0.13.1 to support latest pytorch

76764e7

[pre-commit.ci] auto fixes from pre-commit.com hooks

05c15ba

for more information, see https://pre-commit.ci

krshrimali commented Aug 10, 2022

View reviewed changes

requirements/datatype_text.txt Outdated Show resolved Hide resolved

krshrimali added 2 commits August 10, 2022 09:22

Update requirements/datatype_text.txt

2fc3d82

dont allow 1.7.0 until issues are fixed

18e7144

krshrimali mentioned this pull request Aug 10, 2022

CI: Don't move to Lightning Trainer / PL 1.7.0 (temporary fix) #1415

Closed

8 tasks

krshrimali and others added 3 commits August 10, 2022 15:15

Requirements: go back to PL 1.7

74e7345

CI Fix: manually setattr for collate_fn after DataLoader is initialized

f04a274

[pre-commit.ci] auto fixes from pre-commit.com hooks

cdbaefd

for more information, see https://pre-commit.ci

carmocca reviewed Aug 10, 2022

View reviewed changes

krshrimali mentioned this pull request Aug 11, 2022

flash_examples object_detection trainer.predict not work #1416

Closed

krshrimali added 2 commits August 12, 2022 09:50

install effdet from latest commit hash, fix object_detection

b410d31

Force install effdet 0.3.0

a8f57ad

krshrimali and others added 8 commits August 12, 2022 10:27

Remove dataloader_idx from on_train_batch_end hooks

598d978

rename, attempt to support previous PL versions

02826a3

[pre-commit.ci] auto fixes from pre-commit.com hooks

cf1d6e9

for more information, see https://pre-commit.ci

syntax error fix

ea2679e

Merge branch 'flash-ci/functions-to-classes' of github.com:Lightning-…

aa3e00e

…AI/lightning-flash into flash-ci/functions-to-classes

Fix updating collate fn if input_transform is not None (icevision)

4241ac9

Merge branch 'flash-ci/functions-to-classes' of github.com:Lightning-…

d387308

…AI/lightning-flash into flash-ci/functions-to-classes

pep8 fix

802ab2b

krshrimali marked this pull request as ready for review August 12, 2022 10:41

ethanwharris reviewed Aug 12, 2022

View reviewed changes

krshrimali and others added 2 commits August 12, 2022 16:47

Revert effdet changes, address reviews

23de0a4

[pre-commit.ci] auto fixes from pre-commit.com hooks

f609ef6

for more information, see https://pre-commit.ci

krshrimali commented Aug 12, 2022

View reviewed changes

krshrimali and others added 3 commits August 12, 2022 16:52

Apply suggestions from code review

7a24767

[pre-commit.ci] auto fixes from pre-commit.com hooks

e58b1d5

for more information, see https://pre-commit.ci

indentation fix

4b32904

krshrimali changed the title ~~Fix Flash CLI: in case of converting functions to classes + JIT tracing tests~~ Lightning Trainer/PyTorch Lightning 1.7.0 support + CI Fixes (JIT Tracing and Functions to Classes conversion) Aug 12, 2022

krshrimali requested a review from ethanwharris August 12, 2022 13:07

krshrimali commented Aug 12, 2022

View reviewed changes

.azure-pipelines/testing-template.yml Outdated Show resolved Hide resolved

Update .azure-pipelines/testing-template.yml

2c89478

krshrimali commented Aug 12, 2022

View reviewed changes

.azure-pipelines/testing-template.yml Outdated Show resolved Hide resolved

Update .azure-pipelines/testing-template.yml

ccc7e24

ethanwharris approved these changes Aug 12, 2022

View reviewed changes

krshrimali added 2 commits August 12, 2022 18:46

Add CHANGELOG entries

ef9b0c0

Merge branch 'flash-ci/functions-to-classes' of github.com:Lightning-…

8e03fd7

…AI/lightning-flash into flash-ci/functions-to-classes

ethanwharris merged commit 2657315 into master Aug 12, 2022

ethanwharris deleted the flash-ci/functions-to-classes branch August 12, 2022 13:59

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

krshrimali commented Aug 4, 2022 •

edited

Loading

codecov bot commented Aug 4, 2022 •

edited

Loading

carmocca Aug 10, 2022

krshrimali Aug 12, 2022

krshrimali commented Aug 12, 2022

ethanwharris left a comment •

edited

Loading

ethanwharris Aug 12, 2022

otaj Aug 12, 2022

krshrimali Aug 12, 2022

ethanwharris Aug 12, 2022

krshrimali Aug 12, 2022

krshrimali Aug 12, 2022

krshrimali commented Aug 12, 2022

ethanwharris left a comment

ethanwharris Aug 12, 2022

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

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

Conversation

krshrimali commented Aug 4, 2022 • edited Loading

What does this PR do?

Before submitting

PR review

Did you have fun?

codecov bot commented Aug 4, 2022 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

krshrimali commented Aug 12, 2022

ethanwharris left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

krshrimali commented Aug 12, 2022

ethanwharris left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

krshrimali commented Aug 4, 2022 •

edited

Loading

codecov bot commented Aug 4, 2022 •

edited

Loading

ethanwharris left a comment •

edited

Loading