Avoid importing tensorflow when importing evaluate by NouamaneTazi · Pull Request #135 · huggingface/evaluate

NouamaneTazi · 2022-06-12T12:20:32Z

Avoid loading TFPreTrainedModel when it's not used because tensorflow allocates all GPU memory

Fixes the issue:

import evaluate # this would allocate all GPU memory

Note: also made a fix on pipeline to work as expected huggingface/transformers#17684

HuggingFaceDocBuilderDev · 2022-06-12T12:23:20Z

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

ola13

Looks good, thanks @NouamaneTazi! How did you test the change?

lhoestq

I think this won't work for users-defined models that inherit from PreTrainedModel.

I would simply import PreTrainedModel /TFPreTrainedModel directly here inside compute instead of importing it at the top of the file.

Calling compute is expensive anyway, so importing tensorflow is ok I think.

(alternatively you can check if tensorflow has already been imported by checking "tensorflow" in sys.modules AFAIK - if it's not been imported you don't need to import TFPreTrainedModel and check if the model inherits from TFPreTrainedModel)

NouamaneTazi · 2022-06-13T08:21:26Z

Using the fix in this PR, importing evaluate should no longer allocate all the GPU memory.

And by using the PR huggingface/transformers#17684 as well, e.compute() should no longer allocate all the GPU memory neither.

from evaluate import evaluator
from datasets import Dataset, load_dataset

e = evaluator("text-classification")
data = Dataset.from_dict(load_dataset("imdb")["test"][:2])

# testing that nothing breaks
# from transformers import TFBertModel, BertModel
# tfmodel = TFBertModel.from_pretrained("julien-c/bert-xsmall-dummy")
# model = BertModel.from_pretrained("julien-c/bert-xsmall-dummy")

results = e.compute(
    model_or_pipeline="huggingface/prunebert-base-uncased-6-finepruned-w-distil-mnli",
    data=data,
    metric="accuracy",
    input_column="text",
    label_column="label",
    label_mapping={"LABEL_0": 0.0, "LABEL_1": 1.0},
    strategy="bootstrap",
    n_resamples=10,
    random_state=0
)

NouamaneTazi · 2022-06-13T08:33:16Z

I'm not sure if "tensorflow" in sys.modules is the way to do it. Because if the user uses pipeline, that means they don't know in advance what's the model type right?

And I agree about the user-defined models @lhoestq

NouamaneTazi · 2022-06-13T08:44:59Z

Testing script for user-defined model

import torch
from datasets import Dataset, load_dataset
from transformers import BertConfig, BertTokenizer, PreTrainedModel

from evaluate import evaluator


e = evaluator("text-classification")
data = Dataset.from_dict(load_dataset("imdb")["test"][:2])


class CustomModel(PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
    def forward(self, *args, **kwargs):
        return {'logits': torch.zeros(1, 1, 1)}


model_or_pipeline = CustomModel(BertConfig.from_pretrained("julien-c/bert-xsmall-dummy"))

results = e.compute(
    model_or_pipeline=model_or_pipeline,
    tokenizer=BertTokenizer.from_pretrained("julien-c/bert-xsmall-dummy"),
    data=data,
    metric="accuracy",
    input_column="text",
    label_column="label",
    label_mapping={"LABEL_0": 0.0, "LABEL_1": 1.0},
    strategy="bootstrap",
    n_resamples=10,
    random_state=0,
)

lhoestq · 2022-06-13T09:39:57Z

+        if isinstance(model_or_pipeline, str) or (
+            hasattr(model_or_pipeline, "__class__")
+            and any(
+                cls_name in [parent_cls.__name__ for parent_cls in model_or_pipeline.__class__.__mro__]
+                for cls_name in ["PreTrainedModel", "TFPreTrainedModel"]
+            )


I'm not sure if "tensorflow" in sys.modules is the way to do it. Because if the user uses pipeline, that means they don't know in advance what's the model type right?

If model_or_pipeline is a pipeline, you don't need to check if it inherits from TFPreTrainedModel, so you don't need to import TFPreTrainedModel ;)

I'll let @sgugger give his opinion, but I think this is reasonable:

Suggested change

if isinstance(model_or_pipeline, str) or (

hasattr(model_or_pipeline, "__class__")

and any(

cls_name in [parent_cls.__name__ for parent_cls in model_or_pipeline.__class__.__mro__]

for cls_name in ["PreTrainedModel", "TFPreTrainedModel"]

)

import transformers

if "tensorflow" in sys.modules:

# Check if the model if a TF model only if TF has already been imported.

# Indeed loading `TFPreTrainedModel` may import tensorflow unnecessarily.

transformers_pretrained_model_classes = (transformers.PreTrainedModel, transformers.TFPreTrainedModel)

else:

transformers_pretrained_model_classes = (transformers.PreTrainedModel)

if (

isinstance(model_or_pipeline, transformers_pretrained_model_classes)

or isinstance(model_or_pipeline, str)

After testing this, it seems that after this line from datasets import Dataset, load_dataset, "tensorflow" in sys.modules becomes True. And so it doesn't help with the problem.
Is it just my python?

Yea it's an issue with datasets, let me fix it

Actually it comes from huggingface_hub - let me open a PR

It's already fixed on the main branch of huggingface-hub. I also added a tests here for the future: huggingface/huggingface_hub#904 and huggingface/datasets#4482

This suggestion might avoid the load yes. The current test as it is useless as no model is ever just a PreTrainedModel or a TFPreTrainedModel. They are always subclasses of those.

The current test as it is useless as no model is ever just a PreTrainedModel or a TFPreTrainedModel. They are always subclasses of those.

indeed, using isinstance is required here and it's fine to import PreTrainedModel (and also TFPreTrainedModel if tensorflow has been imported by the user) - see my suggestion above

I said: "This suggestion might avoid the load yes." ;-)
The rest of the comment is on the existing code.

This is actually not enough to not import TF, since importing pipeline (from transformers.pipelines) does import TF. Is this expected ?

I wouldn't spend too much time in evaluate trying to not load TF at this point. Users can still provide USE_TF=0 to not import TF in transformers.

It's maybe a better solution to lazy import the evaluator module, so that transformers is not imported when evaluate is imported

Yes pipeline imports TensorFlow if available to load the default model associated with each pipeline. I would recommend not importing pipeline before the time it's necessary.

NouamaneTazi · 2022-06-30T21:06:45Z

        if (
-            isinstance(model_or_pipeline, PreTrainedModel)
-            or isinstance(model_or_pipeline, TFPreTrainedModel)
-            or isinstance(model_or_pipeline, str)
+            isinstance(model_or_pipeline, str)
+            or isinstance(model_or_pipeline, transformers.PreTrainedModel)
+            or isinstance(model_or_pipeline, transformers.TFPreTrainedModel)
        ):


After some more debugging, here's what I found:

import tensorflow # doesn't allocate all GPU memory import transformers # doesn't allocate all GPU memory from transformers import TFPreTrainedModel # allocates all GPU memory

So as long as we avoid using/importing transformers.TFPreTrainedModel, there shouldn't be a problem. Which is why I believe this is the simplest fix

NouamaneTazi · 2022-06-30T21:12:30Z

Again, you can test that this is working by using

from evaluate import evaluator
from datasets import Dataset, load_dataset

e = evaluator("text-classification")
data = Dataset.from_dict(load_dataset("imdb")["test"][:2])

# testing that nothing breaks
# from transformers import TFBertModel, BertModel
# tfmodel = TFBertModel.from_pretrained("julien-c/bert-xsmall-dummy")
# model = BertModel.from_pretrained("julien-c/bert-xsmall-dummy")

results = e.compute(
    model_or_pipeline="huggingface/prunebert-base-uncased-6-finepruned-w-distil-mnli",
    data=data,
    metric="accuracy",
    input_column="text",
    label_column="label",
    label_mapping={"LABEL_0": 0.0, "LABEL_1": 1.0},
    strategy="bootstrap",
    n_resamples=10,
    random_state=0
)

And please note that calling evaluator.compute would call pipeline which would still call a TFPreTrainedModel and allocates all GPU memory. Which is why this PR is necessary

lhoestq

Sounds good to me, indeed it won't call transformers.TFPretrainedModel if the input is a string or a pytorch model :)

lvwerra · 2022-07-06T10:30:49Z

If @sgugger agrees, I think we can merge this.

sgugger · 2022-07-06T12:35:02Z

Double-checked and indeed import pipeline does not allocate GPU memory while import TFPreTrainedModel does (which is a bug on Transformers probably, will dig more). So this should work (but ultimately be rendered unnecessary I hope!)

sgugger · 2022-07-06T14:00:28Z

huggingface/transformers#18044 should fix the fact that importing TFPreTrainedModel takes all GPU memory in the Transformers side.

- avoids importing xxPreTrainedModel when checking instance

NouamaneTazi · 2022-07-15T20:03:59Z

This PR is no longer necessary as the problem was solved from transformers side. But I think it's still better to avoid unnecessary imports, and to check instances in the order PT -> TF (just like it's done in transformers)

lvwerra · 2022-07-18T07:37:41Z

Still a good idea to have it so we don't need to pin a too new transformers version I think.

avoid loading TFPreTrainedModel to check model's type

77ef609

ola13 approved these changes Jun 13, 2022

View reviewed changes

lvwerra requested review from lhoestq and sgugger June 13, 2022 07:54

lhoestq reviewed Jun 13, 2022

View reviewed changes

fix compute() for user-defined models

38174b0

NouamaneTazi requested a review from lhoestq June 13, 2022 08:45

lhoestq reviewed Jun 13, 2022

View reviewed changes

NouamaneTazi changed the title ~~Avoid loading TFPreTrainedModel in evaluator.py~~ Avoid importing tensorflow when importing evaluate Jun 16, 2022

lvwerra requested review from sgugger and removed request for sgugger June 23, 2022 12:45

avoid importing transformers.TFPreTrainedModel unless necessary

d6954f0

NouamaneTazi commented Jun 30, 2022

View reviewed changes

NouamaneTazi requested a review from lhoestq June 30, 2022 21:13

lhoestq approved these changes Jul 4, 2022

View reviewed changes

NouamaneTazi added 2 commits July 15, 2022 21:53

Merge remote-tracking branch 'origin/main' into pr/NouamaneTazi/135

14af63c

remove unused imports and small refactor

4720a26

- avoids importing xxPreTrainedModel when checking instance

lvwerra merged commit 9b6cea3 into huggingface:main Jul 18, 2022

-        if isinstance(model_or_pipeline, str) or (
-            hasattr(model_or_pipeline, "__class__")
-            and any(
-                cls_name in [parent_cls.__name__ for parent_cls in model_or_pipeline.__class__.__mro__]
-                for cls_name in ["PreTrainedModel", "TFPreTrainedModel"]
-            )
+        import transformers
+        if "tensorflow" in sys.modules:
+            # Check if the model if a TF model only if TF has already been imported.
+            # Indeed loading `TFPreTrainedModel` may import tensorflow unnecessarily.
+            transformers_pretrained_model_classes = (transformers.PreTrainedModel, transformers.TFPreTrainedModel)
+        else:
+            transformers_pretrained_model_classes = (transformers.PreTrainedModel)
+        if (
+            isinstance(model_or_pipeline, transformers_pretrained_model_classes)
+            or isinstance(model_or_pipeline, str)

Conversation

NouamaneTazi commented Jun 12, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Jun 12, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ola13 left a comment

Choose a reason for hiding this comment

Uh oh!

lhoestq left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NouamaneTazi commented Jun 13, 2022

Uh oh!

NouamaneTazi commented Jun 13, 2022

Uh oh!

NouamaneTazi commented Jun 13, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lhoestq Jun 13, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lhoestq Jun 30, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NouamaneTazi commented Jun 30, 2022

Uh oh!

lhoestq left a comment

Choose a reason for hiding this comment

Uh oh!

lvwerra commented Jul 6, 2022

Uh oh!

sgugger commented Jul 6, 2022

Uh oh!

sgugger commented Jul 6, 2022

Uh oh!

NouamaneTazi commented Jul 15, 2022

Uh oh!

lvwerra commented Jul 18, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

NouamaneTazi commented Jun 12, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented Jun 12, 2022 •

edited

Loading

lhoestq left a comment •

edited

Loading

lhoestq Jun 13, 2022 •

edited

Loading

lhoestq Jun 30, 2022 •

edited

Loading