Update README.md

README.md

@@ -38,11 +38,15 @@ model='DiT-g/2'                                 # model type
 task_flag="lora_jade_ema_rank64"                # task flag 
 resume=./ckpts/t2i/model/                       # resume checkpoint 
 index_file=dataset/index_v2_json/jade.json      # index file
+results_dir=./log_EXP                           # save root for results
 batch_size=1                                    # training batch size  
+image_size=1024                                 # training image resolution
 grad_accu_steps=2                               # gradient accumulation steps 
+warmup_num_steps=0                              # warm-up steps
+lr=0.0001                                       # learning rate
+ckpt_every=100                                  # create a ckpt every a few steps.
+ckpt_latest_every=2000                          # create a ckpt named `latest.pt` every a few steps. 
 rank=64                                         # rank of lora 
-max_training_steps=2000                         # max training steps  
-lr=0.0001                                       # learning rate 
 
 PYTHONPATH=./ deepspeed hydit/train_large_deepspeed.py \
     --task-flag ${task_flag} \
@@ -58,20 +62,19 @@ PYTHONPATH=./ deepspeed hydit/train_large_deepspeed.py \
     --uncond-p 0.44 \
     --uncond-p-t5 0.44 \
     --index-file ${index_file} \
-    --random-crop \
     --random-flip \
     --batch-size ${batch_size} \
-    --image-size 1024 \
+    --image-size ${image_size} \
     --global-seed 999 \
     --grad-accu-steps ${grad_accu_steps} \
-    --warmup-num-steps 0 \
+    --warmup-num-steps ${warmup_num_steps} \
     --use-flash-attn \
     --use-fp16 \
     --ema-dtype fp32 \
-    --results-dir ./log_EXP \
-    --ckpt-every 100 \
+    --results-dir ${results_dir} \
+    --ckpt-every ${ckpt_every} \
     --max-training-steps ${max_training_steps}\
-    --ckpt-latest-every 2000 \
+    --ckpt-latest-every ${ckpt_latest_every} \
     --log-every 10 \
     --deepspeed \
     --deepspeed-optimizer \