Repository navigation
Add training support for EnCodec #24295
Description
Feature request
Would be cool to add training support for the EnCodec model.
Not entirely sure if we can easily make it compatible with Trainer, so this can be a good second issue I think.
Motivation
…
Your contribution
…
Activity
Swastyy commentedon Jun 17, 2023
Hi @ArthurZucker , I want to try this, please assign it to me. Thanks.
ArthurZucker commentedon Jun 22, 2023
Sure! Feel free to open a PR and ping me
hackyon commentedon Jun 26, 2023
@Swastyy @ArthurZucker Let me know if you are looking for any support. I would also like to help with this if possible. Thanks!
ArthurZucker commentedon Jun 26, 2023
Seems like he did not link a PR, feel free to synch and ping me for any help! Even a draft is good!
rishabbala commentedon Jun 26, 2023
Hi @ArthurZucker can you let me know of the overall changes that have to be made. I see the EnCodec model already implemented in transformers, so to integrate it with Trainer what are the additional requirements?
ArthurZucker commentedon Jun 27, 2023
The idea is mostly to integrate the loss computation for the VQVAE! Trainer might not work as the model does not use attention, but the target should be to have the same loss as the original model !
hackyon commentedon Jun 27, 2023
Thanks Arthur.
I read through the paper (https://arxiv.org/pdf/2210.13438.pdf) and the existing code, and here is my impression on the work breakdown. Does any of this make sense or am I going in a totally wrong direction?
The loss function detailed in the paper (equation 4) is a combination of (1) the reconstruction loss (over frequency and time domains), (2) the discriminative loss (requires the discriminator), and (3) the VQ commitment loss (the quantizer loss).
(1) The reconstruction loss is computed using the original audio as the label, and we basically need to apply certain time and frequency transformations to the input/output and compute the L1/L2 distances between them.
(2) The discriminative loss requires a discriminator. As far as I can tell, this hasn't been ported/implemented yet and we'll need to do it if we wanted to compute the loss as stated in the paper (msstftd.py from facebookresearch). We'll need to hook up the discriminator in the training code somewhere (is there any pretrained discriminator here?). Also, it's unclear to me whether we can train the discriminator and the model/generator at the same time (I'm assuming not, and we'll need to train one at a time).
(3) The VQ commitment loss is from the quantizer. It looks like it's summing up the losses across all the residual steps. Are we supposed to train the quantizer at the same time as the encoder/decoders? Or should we train them at different times?
In addition to the general loss function, the paper introduced a balancer (balancer.py) that weighs the reconstruction, discriminative, and commitment losses differently. We would also need to import the balancer code if we want this special balancer.
ArthurZucker commentedon Jun 27, 2023
Makes sense to me! I think you can focus simply on returning the loss for the modules. The order of training is not that important (when implementing the module wise loss) since you don't need to train (but compare output losses) until you have eveything!
For the discriminator, you can live it in the training file! It should be pretty small and that's usually how we do things 🤗 !
The order of training, on what is frozen when should be in the paper/original codebase, have not looked it up!
hackyon commentedon Jun 27, 2023
I'll attempt to code up (3) VQ commitment loss first then. I'll reach out if I get stuck or run into any issues. Thanks!
hackyon commentedon Jun 28, 2023
I added an initial draft here: 4f697be
Can you take a look and let me know what you think? Thanks
hackyon commentedon Jun 30, 2023
FYI I will be traveling in July, so won't be as available that month.
ArthurZucker commentedon Jun 30, 2023
Sure, would you mind opening a proper PR? Would be easier to test locally and visualize and follow changes!
ZhikangNiu commentedon Aug 26, 2023
So cool, I reproduced the code and release the code. If you have any question, we can solve together. https://github.com/NoFish-528/encodec-pytorch @hackyon @ArthurZucker In this work, I haven't add balancer, it's difficult for me... Hope you can successful
jonflynng commentedon Aug 24, 2024
Hey @ArthurZucker I'm familiar with EnCodec so can take a stab at this if it still seems like a good feature to implement.
ArthurZucker commentedon Aug 27, 2024
Up to you, I think @ylacombe might have some code as they trained the model for parler no? 🤗
ylacombe commentedon Sep 12, 2024
Hey @jonflynng and @ArthurZucker, I haven't actually trained an Encodec yet. Feel free to take a stab at it @jonflynng !
ylacombe commentedon Sep 12, 2024
@hackyon add a PR opened for this #24593 but closed it since. Feel free to start from scratch or to take inspiration from it. Either way, it'd be great if you could take a look to it since @sanchit-gandhi made a few relevant observations!
jonflynng commentedon Sep 16, 2024
@ylacombe yep, I'll start taking a look this week
jonflynng commentedon Oct 4, 2024
@ArthurZucker @ylacombe
I've got a rough first draft, will take a look more in detail over the next week. I was wondering what the ideal API design is for a model like this with multiple moving parts (and models), as currently I've put a few of the loss calculations within the training loop itself but I imagine that's not what we want? https://github.com/huggingface/transformers/pull/33956/files#diff-5d5859fd55a63733e8f3a00ea4be76b7c770bfa5e254d13408ab1185fc565503R186-R318
ArthurZucker commentedon Oct 6, 2024
Hey! I think the best API is to compute the loss outside the modeling. The loss is specific to the model, but maybe specific to training. IMO we should have a separate file like a
loss_encodec.pyin the folder for now, properly separating from the modeling!jonflynng commentedon Oct 20, 2024
Ok, have created that file. In terms of testing, are usually a few integration tests enough or are small scale experiments wanted before putting this into production? I've ran a couple fine-tuning runs, but because EnCodec is already well pre-trained I'm not getting any performance increases when fine-tuning on some audio and speech datasets. The generator loss is stagnant throughout.
ArthurZucker commentedon Oct 22, 2024
We recently added the entire api for this in #34191 ! (individual loss files)
Small scale experiments (that take like max 4min to run) are nice once we know for sure that it "works" and trains properly!
jonflynng commentedon Oct 25, 2024
Ah nice, ok. What are the benefits of this design by the way? In terms of EnCodec, I can make the commitment and reconstruction losses which are computed inside the EnCodecModel class use this new design. But the other losses (feature matching, discriminator, adversarial) are computed outside of the model in the training loop so would they stay as they are in that separate file and then just be imported in when needed?
ArthurZucker commentedon Oct 29, 2024
Yep that's how it would best IMO!
jonflynng commentedon Oct 31, 2024
Ok, have done that. I'll have another look at the small-scale experiment. I was doing it in a Colab instance. What exactly would you like to see, just some wandb logs of the experiment and the code maybe?
ArthurZucker commentedon Nov 25, 2024
Mostly wandb logs with small convergence for example 🤗
ArthurZucker commentedon Oct 20, 2026
Still open for contributions: the last attempt (#33956) was closed because it targeted v4, with an invite to reopen on main. A new PR would add the losses (reconstruction, commitment, and optionally the adversarial/feature-matching terms) to
EncodecModelinsrc/transformers/models/encodec/modeling_encodec.py, ideally with a small convergence run (wandb logs) as ArthurZucker asked above. https://huggingface.co/docs/transformers/main/en/contributing has the PR checklist.