Merge pull request AILab-CVC#32 from KU-AGI/seed_agi

Update Stage 1 Proposal

.vscode/launch.json

@@ -9,10 +9,13 @@
             "type": "python",
             "request": "launch",
             "program": "${file}",
-            // "env": {"CUDA_VISIBLE_DEVICES":"6, 7, 8, 9"},
-            "env": {"CUDA_VISIBLE_DEVICES":"9"},
-            "args" : ["cfg_path=configs/training/stage2/seed_stage2_training.yaml"],
-            // "args" : ["cfg_path=configs/training/stage1/long_caption_test.yaml"],
+            // "env": {"CUDA_VISIBLE_DEVICES":"8, 9"},
+            "env": {"CUDA_VISIBLE_DEVICES":"5"},
+            // "env": {"CUDA_VISIBLE_DEVICES":"0,1,2,3"},
+            // "args" : ["cfg_path=configs/training/stage2/seed_stage2_training.yaml"],
+            // "args" : ["cfg_path=configs/eval/seed_FID.yaml"],
+            // "args" : ["cfg_path=configs/training/stage1/long_caption_training.yaml"],
+            "args" : ["cfg_path=configs/training/stage1/stage_1_training_long_caption_debug.yaml"],
             "console": "integratedTerminal",
             "justMyCode": false,
         },

coco_dataloader.py

@@ -34,7 +34,7 @@ def __init__(self,
             self.karpathy = json.load(f)
 
         self.start_index = start_index
-        self.end_index = end_index
+        self.end_index = None if end_index == "None" else end_index
 
     def __len__(self):
         if self.start_index is not None and self.end_index is not None:

configs/data/dci_llava.yaml

@@ -0,0 +1,4 @@
+META:
+  - ["/ssd0/data/DCI_LLaVA_wds/{00000..00001}.tar", 20000]
+
+CONTAIN_TEXT: True

configs/data/laion_capsfusion.yaml

@@ -0,0 +1,5 @@
+META:
+  - ["/ssd0/data/laion_capsfusion_wds/{00001..01380}_000000.tar", 11152883]
+  # - ["/ssd0/data/laion_capsfusion_wds/{00000..00001}_000000.tar", 20000]
+
+CONTAIN_TEXT: False

configs/data/laion_capsfusion_sc.yaml

@@ -0,0 +1,4 @@
+META:
+  - ["/ssd0/data/laion_capsfusion_sc_wds/{00001..01380}_000000.tar", 11152883]
+
+CONTAIN_TEXT: False

configs/data/laion_capsfusion_val.yaml

@@ -0,0 +1,4 @@
+META:
+  - ["/ssd0/data/laion_capsfusion_wds/{00000..00000}_000000.tar", 10000]
+
+CONTAIN_TEXT: False

configs/eval/seed_FID.yaml

@@ -0,0 +1,60 @@
+cfg_path: ???
+tokenizer_cfg_path: configs/tokenizer/seed_llama_tokenizer_hf.yaml
+transform_cfg_path: configs/transform/clip_transform.yaml
+model_cfg_path: configs/llm/seed_llama_8b.yaml 
+result_file_path: ./logs/seed_FID_not_bypass_codebook
+checkpoint_path:
+  model_path: pretrained/seed_tokenizer/seed_quantizer.pt
+  diffusion_model_path: stabilityai/stable-diffusion-2-1-unclip
+
+resume: False
+load_weight: False
+weight_path: None
+eval: False
+
+dist:
+  n_gpus: 4
+  n_nodes: 1
+
+dataset:
+  val_config:
+    karpathy_file_path: /ssd0/data/coco/annotations/karpathy/dataset_coco_test.json
+    root_dir: /ssd0/data/coco/images/val2014
+  num_workers: 16
+  shuffle: True
+  text_max_length: 128
+
+stage1:
+  init: 'SEED'
+
+stage2:
+  bypass_codebook: False
+  load_diffusion: True
+
+experiment:
+  seed: 0
+  stage: 2
+  local_batch_size: 1024
+  val_batch_size: 16
+  test_split: train
+  max_epochs: 1
+  deterministic: False
+  grad_accumulation: 1
+  check_val_every_n_epoch: 1
+  enable_checkpointing: True
+  log_every_n_steps: 1
+  num_sanity_val_steps: 1
+  num_warmup_steps: 200
+  grad_clip_val: 0.5
+
+optimizer:
+  vit_precision: 'fp16'
+  diffusion_precision: 'fp32'
+  precision: 'bf16'
+  max_lr: 7e-4
+
+hyperparameters:
+  beta_1: 0.9
+  beta_2: 0.999
+  weight_decay: 1e-8
+

configs/seed_training_proj_test.yaml

@@ -30,7 +30,7 @@ dataset:
     type: dalle-vqvae
     hparams:
       resolution: 256
-  gt_text: Trueq
+  gt_text: True
 
 stage1:
   ema_update: False

configs/seed_unified_test_c2f.yaml

@@ -0,0 +1,78 @@
+cfg_path: ???
+tokenizer_cfg_path: configs/tokenizer/seed_llama_tokenizer_hf.yaml
+transform_cfg_path: configs/transform/clip_transform.yaml
+model_cfg_path: configs/llm/seed_llama_8b.yaml 
+result_file_path: ./logs/noexp
+checkpoint_path:
+  model_path: pretrained/seed_tokenizer/seed_quantizer.pt
+  diffusion_model_path: stabilityai/stable-diffusion-2-1-unclip
+
+resume: False
+load_weight: False
+weight_path: None
+eval: False
+
+dist:
+  n_gpus: 1
+  n_nodes: 1
+
+dataset:
+  train_config:
+    dataset_configs: ['configs/data/laion_capsfusion_sc.yaml']
+    weights: [1]
+    shardshuffle: 100
+    resampled: True
+    world_size: 1
+    one_epoch_data_size: 1000000
+  val_config:
+    karpathy_file_path: /ssd0/data/coco/annotations/karpathy/dataset_coco_test.json
+    root_dir: /ssd0/data/coco/images/val2014
+    start_index: 0
+    end_index: -1
+  num_workers: 1
+  shuffle: True
+  text_max_length: 128
+
+stage1:
+  init: 'SEED'
+
+stage2:
+  bypass_codebook: True
+  load_diffusion: False
+  train_unet: False
+
+experiment:
+  seed: 0
+  stage: 1
+  local_batch_size: 2
+  val_batch_size: 8
+  test_split: train
+  max_epochs: 40
+  deterministic: False
+  grad_accumulation: 8
+#  check_val_every_n_epoch: 1
+  val_check_interval: 400
+  enable_checkpointing: True
+  log_every_n_steps: 1
+  num_sanity_val_steps: 0
+  num_warmup_steps: 200
+  recon_loss_weight: 1.0
+  sds_loss_weight: 0.1
+  clip_loss_weight: 1.0
+  use_sds_loss_schedule: True
+  cross_annealing: True
+  num_positive_samples: 4
+  min_pos_weight: 0.3
+
+optimizer:
+  vit_precision: 'fp16'
+  diffusion_precision: 'fp32'
+  precision: 'bf16'
+  max_lr: 7e-4
+  grad_clip_val: 0.5
+
+hyperparameters:
+  beta_1: 0.9
+  beta_2: 0.999
+  weight_decay: 1e-8
+

configs/seed_unified_test_sds.yaml

@@ -2,7 +2,7 @@ cfg_path: ???
 tokenizer_cfg_path: configs/tokenizer/seed_llama_tokenizer_hf.yaml
 transform_cfg_path: configs/transform/clip_transform.yaml
 model_cfg_path: configs/llm/seed_llama_8b.yaml 
-result_file_path: ./logs/sds_coco
+result_file_path: ./logs/sds_coco2
 checkpoint_path:
   model_path: pretrained/seed_tokenizer/seed_quantizer.pt
   diffusion_model_path: stabilityai/stable-diffusion-2-1-unclip
@@ -18,52 +18,56 @@ dist:
 
 dataset:
   train_config:
-    dataset_configs: ['configs/data/cc15m.yaml']
-    weights: [1]
+    dataset_configs: ['configs/data/cc15m.yaml', 'configs/data/laion-coco.yaml', 'configs/data/mscoco.yaml']
+    weights: [1, 8, 1]
     shardshuffle: 100
     resampled: True
     world_size: 1
-    one_epoch_data_size: 3000000
+    one_epoch_data_size: 1000000
   val_config:
     karpathy_file_path: /ssd0/data/coco/annotations/karpathy/dataset_coco_test.json
     root_dir: /ssd0/data/coco/images/val2014
     start_index: 0
     end_index: 256
   num_workers: 16
   shuffle: True
+  text_max_length: 128
 
 stage1:
-  init: 'BLIP-2'
+  init: 'SEED'
 
 stage2:
   bypass_codebook: True
   load_diffusion: True
-  train_unet: False
-  use_clip_loss: True
+  train_unet: True
 
 experiment:
   seed: 0
   stage: 2
-  local_batch_size: 32
-  val_batch_size: 4
+  local_batch_size: 128
+  val_batch_size: 8
   test_split: train
-  max_epochs: 5
+  max_epochs: 40
   deterministic: False
   grad_accumulation: 8
-  check_val_every_n_epoch: 1
+#  check_val_every_n_epoch: 1
+  val_check_interval: 400
   enable_checkpointing: True
   log_every_n_steps: 1
   num_sanity_val_steps: 1
-  num_warmup_steps: 50
-  grad_clip_val: 1
-  val_check_interval: 200
-
+  num_warmup_steps: 200
+  recon_loss_weight: 1.0
+  sds_loss_weight: 0.1
+  clip_loss_weight: 1.0
+  use_sds_loss_schedule: True
+  cross_annealing: False
 
 optimizer:
   vit_precision: 'fp16'
-  diffusion_precision: 'fp16'
+  diffusion_precision: 'fp32'
   precision: 'bf16'
-  max_lr: 1e-4
+  max_lr: 7e-4
+  grad_clip_val: 0.5
 
 hyperparameters:
   beta_1: 0.9

configs/seed_unified_test_sds_debug.yaml

@@ -0,0 +1,75 @@
+cfg_path: ???
+tokenizer_cfg_path: configs/tokenizer/seed_llama_tokenizer_hf.yaml
+transform_cfg_path: configs/transform/clip_transform.yaml
+model_cfg_path: configs/llm/seed_llama_8b.yaml 
+result_file_path: ./logs/noexp
+checkpoint_path:
+  model_path: pretrained/seed_tokenizer/seed_quantizer.pt
+  diffusion_model_path: stabilityai/stable-diffusion-2-1-unclip
+
+resume: False
+load_weight: False
+weight_path: None
+eval: False
+
+dist:
+  n_gpus: 1
+  n_nodes: 1
+
+dataset:
+  train_config:
+    dataset_configs: ['configs/data/cc15m.yaml', 'configs/data/laion-coco.yaml', 'configs/data/mscoco.yaml']
+    weights: [1, 8, 1]
+    shardshuffle: 100
+    resampled: True
+    world_size: 1
+    one_epoch_data_size: 1000000
+  val_config:
+    karpathy_file_path: /ssd0/data/coco/annotations/karpathy/dataset_coco_test.json
+    root_dir: /ssd0/data/coco/images/val2014
+    start_index: 0
+    end_index: 256
+  num_workers: 1
+  shuffle: True
+  text_max_length: 128
+
+stage1:
+  init: 'SEED'
+
+stage2:
+  bypass_codebook: True
+  load_diffusion: True
+  train_unet: True
+
+experiment:
+  seed: 0
+  stage: 2
+  local_batch_size: 1
+  val_batch_size: 1
+  test_split: train
+  max_epochs: 40
+  deterministic: True
+  grad_accumulation: 1
+#  check_val_every_n_epoch: 1
+  val_check_interval: 400
+  enable_checkpointing: True
+  log_every_n_steps: 1
+  num_sanity_val_steps: 0
+  num_warmup_steps: 0
+  recon_loss_weight: 1.0
+  sds_loss_weight: 0.1
+  clip_loss_weight: 1.0
+  use_sds_loss_schedule: False
+
+optimizer:
+  vit_precision: 'fp16'
+  diffusion_precision: 'fp32'
+  precision: 'bf16'
+  max_lr: 7e-4
+  grad_clip_val: 0.5
+
+hyperparameters:
+  beta_1: 0.9
+  beta_2: 0.999
+  weight_decay: 1e-8
+