Skip to content

Optimization issue/suggestion for "train_d3po.py" #21

Description

@AHHHZ975

Hi @yk7333 @ADKoishi ,

First, thanks for the amazing work and for making it publicly available! I took a close look at the "train_d3po.py" script and ran it, and I wanted to share a quick observation and check if my understanding of the training process is on point.

While going through the code, I noticed a few inefficiencies, the biggest one is that latents and log probabilities are calculated twice: once during image generation (sampling or forward path) and again during fine-tuning (backward path). Since both calls are identical, it’s causing some unnecessary extra computation.

To fix this, I refactored the code to call the diffusion model just once, and it helped improve efficiency. In my tests, this change dropped the iteration time from around 32 seconds to 23 seconds on my GPU. Everything seems to be running smoothly now, but I wanted to check with you to see if you were already aware of this or if there was a specific reason you implemented it this way. Any insights on this would be appreciated.

Thanks,
Amir

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions