Skip to content

Latest commit

 

History

History
110 lines (67 loc) · 2.36 KB

README.md

File metadata and controls

110 lines (67 loc) · 2.36 KB
Raw

TorchDiffusion

One Diffusion model implementation base on LibTorch, The LibTorch version is almost transfer from my pytorch implementation https://github.com/GlassyWing/ddpm

Supported Sampler

  • DDPM
  • DDIM
  • RectifiedFlow

Install

on linux

mkdir build
cd build
cmake  -DCMAKE_PREFIX_PATH="/path/to/libtorch;/path/to/opencv" ..

on windows (assume the Visual Studio 2019/2022 has installed)

mkdir build
cd build
cmd /C "C:\Program Files\Microsoft Visual Studio\2022\Community\Common7\IDE\CommonExtensions\Microsoft\CMake\CMake\bin\cmake"  -DCMAKE_PREFIX_PATH="/path/to/libtorch;/path/to/opencv" ..

The project will created under the directory build.

Build

Using the command to create one Release version.

cd build
cmake --build . --config Release

The runnable program will be build, the location depends on your OS, with windows, it's build/Release/diffusion.exe.

Usage

Train

./diffusion --dataset <dataset_dir> -t ddim -s 4

The default will create one experiments folder used to save the checkpoints and log images.

see more with ./diffusion -h

Inference

To generate one 4x4 demo image by this way:

./diffusion --mode test --weight_path <path/to/weight_path> -t ddim -s 8

Generate Example

After 5 days (140k images):

It's trained with this strategy:

Epochs Approximate Batch Size Batch Size accumulation_steps
0-40 64 64 1
40-80 128 64 2
80-120 256 64 4
120-160 512 64 8
160-200 1024 64 16

And this cmd:

.\diffusion.exe -p <path/to/last_checkpoint> -t 'ddim' -s 4 --amp true -b 64 --accum <accumulation_steps>

Tips: You can still accumulate the size of the Batch_size to get better results

QA

  1. How long does it take to train to see considerable results?

    About 30min on 3090.

  2. Memory usage?

    Image Size: 128 x 128

    Batch Size: 32

    Memory: 3GB

    Vedio Memroy: 13GB

  3. Performance compare to pytorch?

    About 10-15% speedups.