Guidance on Training the EnergyDiff (Calibrated) Model with Custom Data

[yusufM03]

Hi,

Great work on this project!

I want to test the approach described in the paper, but I can't find the proposed model. Could you kindly guide me on how to train the model using my custom data?

Thank you!

[sentient-codebot]

Hi @yusufM03

To use the model on your own data, there are a few things you need to do.

Prepare the Data

The most important thing is to prepare your data in the needed format. You need to have a data loader which is an iterable of a tuple (profile, condition). The profile needs to be a tensor of shape ["batch channel sequence"]. The condition can be empty (None).

for batch in your_dataloader:
    profile, conditions = batch

Choose the Configuration

Check the src/energydiff/utils/configuration.py, there you can find how to define the hyperparameters as you like.

Example of Complete Pipeline

Check scripts/train_ddpm_pl.py. Basically, the procedure is

1. Process arguments and configurations
2. Prepare the dataset (data loader)
3. Define the PytorchLightning module pl_model = energydiff.diffusion.PLDiffusion1D(..), which is a wrapper for all the training and validation you need.
4. Define a trainer for the PytorchLightning module trainer = create_pl_trainer(config, run_id). Here you can customize as you like. For example, you can choose not to use wandb.
5. Initiate the training trainer.fit(pl_model).

You can also check another repository OpenSynth (the develop branch). You can find some notebooks and tutorials of my model there.

Calibration

The marginal calibration is independent of the generative models you use. Once you generated some synthetic data (with EnergyDiff or whatever other models), you can use the function calibrate from energydiff.marginal_calibrate. It is simply

calibrated_data = calibrate(
    target: Float[Tensor, 'batch, channel, length'] = some_training_samples,
    source: Float[Tensor, 'batch, channel, length'] = your_generated_data,
)