Skip to content

[Pipeline] avoid importing tensorflow if not used#17684

Closed
NouamaneTazi wants to merge 2 commits intohuggingface:mainfrom
NouamaneTazi:avoid-loading-tf
Closed

[Pipeline] avoid importing tensorflow if not used#17684
NouamaneTazi wants to merge 2 commits intohuggingface:mainfrom
NouamaneTazi:avoid-loading-tf

Conversation

@NouamaneTazi
Copy link
Copy Markdown
Member

What does this PR do?

Avoids loading unnecessary modules by pipelines.base.infer_framework_load_model() which could create some unexpected behaviour like tensorflow allocating all GPU memory.
@sgugger @LysandreJik

Before:

from transformers import pipeline
pipeline("text-classification")
# This would try importing `TFDistilBertForSequenceClassification`if both tensorflow and pytorch
# are available, and tensorflow would allocate all GPU memory, even if we expect to use
# the pytorch model

After:

from transformers import pipeline
pipeline("text-classification")
# Only `DistilBertForSequenceClassification` is imported, and tensorflow is not called

@HuggingFaceDocBuilderDev
Copy link
Copy Markdown

HuggingFaceDocBuilderDev commented Jun 12, 2022

The documentation is not available anymore as the PR was closed or merged.

@LysandreJik LysandreJik requested a review from sgugger June 13, 2022 07:20
Copy link
Copy Markdown
Collaborator

@sgugger sgugger left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks generally ok to me, but would like to see what @Narsil thinks since pipelines are his domain :-)
I can confirm that creating a pipeline(xxx) now imports TensorFlow even if the model is in PyTorch at the end. I don't think it was the case before (when I was writing the course so a while ago 😅), but maybe I misremember.

@NouamaneTazi NouamaneTazi changed the title [Pipeline] avoid loading tensorflow if not used [Pipeline] avoid importing tensorflow if not used Jun 16, 2022
@github-actions
Copy link
Copy Markdown
Contributor

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

@NouamaneTazi
Copy link
Copy Markdown
Member Author

Was this issue resolved in another PR already? @sgugger

@sgugger
Copy link
Copy Markdown
Collaborator

sgugger commented Jul 13, 2022

The fact that TensorFlow takes all GPU memory has been fixed in #18044

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants