[Utils] Adds `store()` and `restore()` methods to EMAModel #2302

sayakpaul · 2023-02-09T11:13:35Z

This PR adds the store() and restore() methods to our EMAModel class. The reason is beautifully explained by @patil-suraj here.

An end-to-end example of how to use these methods in practice is here.

It also un-blocks #2157

HuggingFaceDocBuilderDev · 2023-02-09T11:18:58Z

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

patrickvonplaten

Ok for me! Let's just be a bit careful when adding this functionality to the training scripts to not blow up GPU memory

src/diffusers/training_utils.py

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

sayakpaul · 2023-02-13T14:45:46Z

For the code quality tests, I am running make style and make quality. It picks up a script unrelated to this PR:

src/diffusers/training_utils.py

My formatting virtual environment is up-to-date as well:

black: 23.1.0
isort: 5.12.0
ruff: 0.0.244

sayakpaul · 2023-02-14T08:15:36Z

Still the same issue: #2302 (comment)

patrickvonplaten · 2023-02-14T16:27:45Z

Fixed the make quality issue. @sayakpaul, your hf-doc-builder library might be outdated.

We also use https://github.com/huggingface/doc-builder to improve the style of our code (think it makes sure comments are not too long.

You can install it with:

pip install hf-doc-builder

williamberman · 2023-02-15T17:55:56Z

src/diffusers/training_utils.py

+        parameters = list(parameters)
+        self.collected_params = [param.detach().cpu().clone() for param in parameters]


can we just have the caller do the detach.cpu.clone and save collected_params in a variable themselves? This sort of adhoc storage just between two methods is a bit of an OO antipattern imo

Another thing is I think the initial list() is also unnecessary. If something is an Iterable, I think we should be able to directly iterate over it?

Yes, that is true. The detach code was after, hence.

williamberman · 2023-02-15T17:59:14Z

src/diffusers/training_utils.py

+        if self.collected_params is None:
+            raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights " "to `restore()`")


If we do https://github.com/huggingface/diffusers/pull/2302/files#r1107495110, then this method wouldn't need to reference self and it wouldn't have this error condition. This would mean maybe it would look better as a function, and then store and restore could be free-standing functions not on a class.

WDYT?

Since the EMAModel class majorly going to be used in the examples, which are meant to be easier to work with, I like this design approach better. So, with that in mind, for #2302 (comment), I think it's probably better if we do it ourselves.

From what I have seen is that the community is fairly eager to customize things as per their needs, but with that, we should aim to provide example scripts that are easy as starting points.

WDYT?

I think we can still make the methods easy to work with without the adhoc variable storage. I.e.

# training utils def store(params): return [x.cpu().detach() for x in params] def restore(params, stored_params): for c_param, param in zip(stored_params, parameters): param.data.copy_(c_param.data) # training script stored = store(params) # do something restore(params, stored)

is still friendly and hackable without the adhoc variable storage

Then it is coming down to a preference now :D

I like the class approach better than what you showed above.

I see your point @williamberman, but IMO from a user's perspective, it would be better to just call class methods than having to manually do this.

williamberman

I think these shouldn't be methods on the class but ok with merging if other people feel strongly it should be :)

src/diffusers/training_utils.py

Co-authored-by: Will Berman <wlbberman@gmail.com>

src/diffusers/training_utils.py

patrickvonplaten

I understand @williamberman points here and to be honest the EMAModel class is also not super intuitive to me in general. To me it's actually more of a EMAModelTrainingUtils class.

But from a user point of view that uses the examples training scripts I also think it's more convenient to be able to attach temporarily stored parameters to the EMAModel instead of having some stored list flying around

Maybe some renaming could help here a bit to make it super clear that the EMAModel.store(...) function actually does not store and restore ema parameters but the non_ema parameters?

patrickvonplaten · 2023-02-16T08:46:24Z

Also good to go for me, but would be happy about potentially renaming the functions a bit to make the difference between the actual EMA and the non-EMA weights clearer

patil-suraj

Thanks a lot for the PR. I agree with Patrick regarding naming.
I just left a comment about copying the params CPU in store. Not sure if it's really needed, will need to dig into it a bit. The rest of it looks great!

src/diffusers/training_utils.py

patil-suraj · 2023-02-16T08:55:04Z

src/diffusers/training_utils.py

+        if self.collected_params is None:
+            raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights " "to `restore()`")


I see your point @williamberman, but IMO from a user's perspective, it would be better to just call class methods than having to manually do this.

src/diffusers/training_utils.py

sayakpaul · 2023-02-16T09:18:03Z

I am in agreement regarding the name changes. Applying them now. I will let @patil-suraj play around with the CPU offloading (.cpu()) before then.

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

… into refactor/ema

src/diffusers/training_utils.py

patil-suraj · 2023-02-16T10:36:08Z

We could do this in another PR, but it would be nice to add some tests for ema model. @sayakpaul feel free to work on it if you want.

sayakpaul · 2023-02-16T10:58:21Z

We could do this in another PR, but it would be nice to add some tests for ema model. @sayakpaul feel free to work on it if you want.

I can but in a separate PR.

Co-authored-by: patil-suraj <surajp815@gmail.com>

sayakpaul · 2023-02-16T11:57:03Z

@patil-suraj see if the latest changes are good to go.

patil-suraj

Thanks a lot for updating this quickly. Just left two quick comments, we can remove the temp params from saving/loading state_dict. Should be good to go after that!

src/diffusers/training_utils.py

sayakpaul · 2023-02-16T12:14:30Z

I think the error related to running fix-copies is very likely from an unrelated issue.

patil-suraj · 2023-02-16T13:24:59Z

The EMAModel is in the main init, and hence it's also in the dummy objects (https://github.com/huggingface/diffusers/blob/main/src/diffusers/utils/dummy_pt_objects.py#L618). Whenever we update any object in the main init we need to run make fix-copies.

sayakpaul · 2023-02-16T13:38:12Z

The EMAModel is in the main init, and hence it's also in the dummy objects (https://github.com/huggingface/diffusers/blob/main/src/diffusers/utils/dummy_pt_objects.py#L618). Whenever we update any object in the main init we need to run make fix-copies.

Right. Sorry for missing it out. Should be good now.

patil-suraj · 2023-02-16T14:20:06Z

Thanks! The failing test is unrelated, merging this now:)

…ce#2302) * add store and restore() methods to EMAModel. * Update src/diffusers/training_utils.py Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * make style with doc builder * remove explicit listing. * Apply suggestions from code review Co-authored-by: Will Berman <wlbberman@gmail.com> * Apply suggestions from code review Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * chore: better variable naming. * better treatment of temp_stored_params Co-authored-by: patil-suraj <surajp815@gmail.com> * make style * remove temporary params from earth 🌎 * make fix-copies. --------- Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: Will Berman <wlbberman@gmail.com> Co-authored-by: patil-suraj <surajp815@gmail.com>

add store and restore() methods to EMAModel.

2c887e8

sayakpaul requested review from patil-suraj and patrickvonplaten February 9, 2023 11:13

patrickvonplaten approved these changes Feb 13, 2023

View reviewed changes

patrickvonplaten reviewed Feb 13, 2023

View reviewed changes

src/diffusers/training_utils.py Outdated Show resolved Hide resolved

sayakpaul and others added 2 commits February 13, 2023 20:09

Update src/diffusers/training_utils.py

ee5766e

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

Merge branch 'main' into refactor/ema

265586f

patrickvonplaten and others added 2 commits February 13, 2023 19:22

Merge branch 'main' into refactor/ema

2cf589d

Merge branch 'main' into refactor/ema

5f74d2d

patrickvonplaten requested a review from williamberman February 14, 2023 16:24

make style with doc builder

5986943

williamberman reviewed Feb 15, 2023

View reviewed changes

williamberman approved these changes Feb 15, 2023

View reviewed changes

sayakpaul mentioned this pull request Feb 16, 2023

Not copying the model in inference and assigning ema_weights seems incorrect. #2321

Closed

remove explicit listing.

db7f031

williamberman reviewed Feb 16, 2023

View reviewed changes

src/diffusers/training_utils.py Outdated Show resolved Hide resolved

williamberman reviewed Feb 16, 2023

View reviewed changes

src/diffusers/training_utils.py Outdated Show resolved Hide resolved

Apply suggestions from code review

4125190

Co-authored-by: Will Berman <wlbberman@gmail.com>

patrickvonplaten reviewed Feb 16, 2023

View reviewed changes

src/diffusers/training_utils.py Show resolved Hide resolved

patrickvonplaten reviewed Feb 16, 2023

View reviewed changes

src/diffusers/training_utils.py Show resolved Hide resolved

patrickvonplaten reviewed Feb 16, 2023

View reviewed changes

src/diffusers/training_utils.py Outdated Show resolved Hide resolved

patrickvonplaten approved these changes Feb 16, 2023

View reviewed changes

patil-suraj reviewed Feb 16, 2023

View reviewed changes

sayakpaul and others added 3 commits February 16, 2023 14:48

Apply suggestions from code review

e6f2e37

Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>

Merge branch 'refactor/ema' of https://github.com/huggingface/diffusers…

576adf2

… into refactor/ema

chore: better variable naming.

b3a28c0

patil-suraj reviewed Feb 16, 2023

View reviewed changes

src/diffusers/training_utils.py Outdated Show resolved Hide resolved

src/diffusers/training_utils.py Show resolved Hide resolved

sayakpaul and others added 2 commits February 16, 2023 17:24

better treatment of temp_stored_params

635313c

Co-authored-by: patil-suraj <surajp815@gmail.com>

make style

9adac8d

patil-suraj approved these changes Feb 16, 2023

View reviewed changes

src/diffusers/training_utils.py Outdated Show resolved Hide resolved

src/diffusers/training_utils.py Outdated Show resolved Hide resolved

remove temporary params from earth 🌎

a2e4477

sayakpaul mentioned this pull request Feb 16, 2023

[Utils] Implement tests for EMAModel #2377

Closed

Merge branch 'main' into refactor/ema

6bc48ef

make fix-copies.

5873864

patil-suraj merged commit 6eaebe8 into main Feb 16, 2023

patil-suraj deleted the refactor/ema branch February 16, 2023 14:20

sayakpaul mentioned this pull request Feb 16, 2023

[WIP] Sample images when checkpointing. #2157

Closed

		parameters = list(parameters)
		self.collected_params = [param.detach().cpu().clone() for param in parameters]

		if self.collected_params is None:
		raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights " "to `restore()`")

[Utils] Adds store() and restore() methods to EMAModel #2302

[Utils] Adds store() and restore() methods to EMAModel #2302

Uh oh!

Conversation

sayakpaul commented Feb 9, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Feb 9, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sayakpaul commented Feb 13, 2023

Uh oh!

sayakpaul commented Feb 14, 2023

Uh oh!

patrickvonplaten commented Feb 14, 2023

Uh oh!

williamberman Feb 15, 2023

Choose a reason for hiding this comment

Uh oh!

williamberman Feb 16, 2023

Choose a reason for hiding this comment

Uh oh!

sayakpaul Feb 16, 2023

Choose a reason for hiding this comment

Uh oh!

williamberman Feb 15, 2023

Choose a reason for hiding this comment

Uh oh!

sayakpaul Feb 16, 2023

Choose a reason for hiding this comment

Uh oh!

williamberman Feb 16, 2023 • 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!

sayakpaul Feb 16, 2023 • 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!

patil-suraj Feb 16, 2023

Choose a reason for hiding this comment

Uh oh!

williamberman left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten commented Feb 16, 2023

Uh oh!

patil-suraj left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

patil-suraj Feb 16, 2023

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sayakpaul commented Feb 16, 2023

Uh oh!

Uh oh!

Uh oh!

patil-suraj commented Feb 16, 2023

Uh oh!

sayakpaul commented Feb 16, 2023

Uh oh!

sayakpaul commented Feb 16, 2023

Uh oh!

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

Uh oh!

[Utils] Adds `store()` and `restore()` methods to EMAModel #2302

[Utils] Adds `store()` and `restore()` methods to EMAModel #2302

sayakpaul commented Feb 9, 2023 •

edited

Loading

HuggingFaceDocBuilderDev commented Feb 9, 2023 •

edited

Loading

williamberman Feb 16, 2023 •

edited

Loading

sayakpaul Feb 16, 2023 •

edited

Loading

patil-suraj left a comment •

edited

Loading