Create config.yaml

config.yaml

@@ -0,0 +1,76 @@
+dataset: point_maze
+data_path: /data/user_data/mchen5/dino_wm_dataset/point_maze
+pretrained_model_path: /data/hf_cache/hub/models--maxin-cn--Latte/snapshots/fc35c8c44c38e0b446158ace47e7de642066f0d9/vae
+n_rollout: null
+split_ratio: 0.9
+validation_size: 32
+use_wandb: true
+results_dir: ./results
+pretrained: null
+model: Latte-S/2
+use_dynamic_tokenizer: true
+num_slots: 256
+slot_dim: 16
+dec_hidden_size: 384
+dec_depth: 12
+dec_num_heads: 6
+num_frames: 4
+n_context_frames: 1
+scheduling_mode: full_sequence
+image_size: 256
+num_sampling_steps: 250
+frame_interval: 1
+fixed_spatial: false
+attention_bias: true
+learn_sigma: false
+diffusion_mode: diffusion_forcing
+extras: 1
+use_action: true
+action_dim: 2
+normalize_action: true
+save_ceph: false
+learning_rate: 0.0001
+clip_max_norm: 0.1
+start_clip_iter: 20000
+local_batch_size: 16
+max_train_steps: 1000000
+global_seed: 3407
+num_workers: 4
+log_every: 10
+val_every_n_steps: 1000
+lr_warmup_steps: 0
+resume_from_checkpoint: null
+gradient_accumulation_steps: 1
+num_classes: null
+use_compile: false
+mixed_precision: 32
+gradient_checkpointing: true
+sequential_loss_grad: true
+callbacks:
+  model_checkpoint_across_timesteps:
+    every_n_train_steps: 10000
+    save_top_k: -1
+    save_on_train_epoch_end: true
+    save_weights_only: false
+    filename: '{epoch}-{step}'
+  model_checkpoint_latest:
+    every_n_train_steps: 1000
+    save_top_k: 1
+    save_on_train_epoch_end: false
+    save_weights_only: false
+    filename: latest-{epoch}-{step}
+  metrics_logger:
+    log_every_n_train_steps: 10000
+    save_videos: true
+    evaluate: true
+    max_batchsize: 4
+    buffer_size: 32
+    i3d_model_path: checkpoints/i3d/i3d_torchscript.pt
+    inference_n_slots:
+    - 64
+    - 128
+    - 256
+train: true
+validate: false
+exp_name: tok2_comp=False_mp=32_gc=True
+run_id: 6002143-0