google / google/flax

How to optimize multi model’s parameter in one optimizer ?

Open
#3,524 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Is there something similar in the flax framework? Or how can I accomplish this use case using flax?

![Screenshot 2023-11-30 at 2 34 21 PM](https://github.com/google/flax/assets/87750395/d30ebacc-d415-4cff-9d79-b9c046970f4f)

One straightforward method combines all these models into a single model (class) inherited from nn.Module.
For a particular use case, this approach didn't work. Is there any other way to accomplish this task in flax using optax?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.