• Config G.
• Post-hoc EMA.

git clone https://github.com/YichengDWu/tinyedm.git cd tinyedm && pip install . 

python experiments/train.py --config-name=mnist python experiments/train.py --config-name=cifar10

To download the ImageNet dataset, follow these steps:

1. Visit the ImageNet website: http://www.image-net.org/
2. Register for an account and request access for the dataset.
3. Once approved, follow the instructions provided by ImageNet to download the dataset.

After downloading the ImageNet dataset, extract the files to a directory. When running the feature extraction script, use the --data-dir option to specify the path to this directory.

For example:

python src/tinyedm/datamodules/extract_latents.py --data-dir ./datasets/imagenet/train --out-dir ./datasets/imagenet/latents/train

python src/tinyedm/generate.py \ --ckpt_path /path/to/checkpoint.ckpt \ --load_ema \ --output_dir /path/to/output \ --num_samples 50000 \ --image_size 32 \ --num_classes 10 \ --batch_size 128 \ --num_workers 16 \ --num_steps 32

1. Using FP16 mixed precision training on the CIFAR-10 dataset sometimes leads to overflow, so we have adopted bf16 mixed precision, which may result in a loss of accuracy for the model.
2. For the scale factors of skip connections, this implementation uses a small network to learn them, inspired by ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection . The experiment shows that this improves the results.
3. The use of multi-task learning in the paper did not observe any improvement, or it may be more effective in long-term training. However, I do not have the compute power to verify this.