Training AutoencoderRAE

This example trains the decoder of AutoencoderRAE (stage-1 style), while keeping the representation encoder frozen.

It follows the same high-level training recipe as the official RAE stage-1 setup:

frozen encoder
train decoder
pixel reconstruction loss
optional encoder feature consistency loss

Quickstart

Resume or finetune from pretrained weights

bash

accelerate launch examples/research_projects/autoencoder_rae/train_autoencoder_rae.py \
  --pretrained_model_name_or_path nyu-visionx/RAE-dinov2-wReg-base-ViTXL-n08 \
  --train_data_dir /path/to/imagenet_like_folder \
  --output_dir /tmp/autoencoder-rae \
  --resolution 256 \
  --train_batch_size 8 \
  --learning_rate 1e-4 \
  --num_train_epochs 10 \
  --report_to wandb \
  --reconstruction_loss_type l1 \
  --use_encoder_loss \
  --encoder_loss_weight 0.1

Train from scratch with a pretrained encoder

The following command launches RAE training with "facebook/dinov2-with-registers-base" as the base.

bash

accelerate launch examples/research_projects/autoencoder_rae/train_autoencoder_rae.py \
  --train_data_dir /path/to/imagenet_like_folder \
  --output_dir /tmp/autoencoder-rae \
  --resolution 256 \
  --encoder_type dinov2 \
  --encoder_name_or_path facebook/dinov2-with-registers-base \
  --encoder_input_size 224 \
  --patch_size 16 \
  --image_size 256 \
  --decoder_hidden_size 1152 \
  --decoder_num_hidden_layers 28 \
  --decoder_num_attention_heads 16 \
  --decoder_intermediate_size 4096 \
  --train_batch_size 8 \
  --learning_rate 1e-4 \
  --num_train_epochs 10 \
  --report_to wandb \
  --reconstruction_loss_type l1 \
  --use_encoder_loss \
  --encoder_loss_weight 0.1

Note: stage-1 reconstruction loss assumes matching target/output spatial size, so --resolution must equal --image_size.

Dataset format is expected to be ImageFolder-compatible:

text

train_data_dir/
  class_a/
    img_0001.jpg
  class_b/
    img_0002.jpg