Update README.md

README.md

@@ -53,7 +53,7 @@ You can change aggregate strategies using the `--agg_strategy` flag, such as `sa
 
 ## Train Step 1 ##
 ```
-accelerate launch --config_file=./accelerate_configs/deepspeed_zero2.yaml run_wsi.py --learning_rate 1e-4 --max_steps 10000 --warmup_steps 1000\
+accelerate launch --config_file=./accelerate_configs/deepspeed_zero2.yaml run_wsi.py --learning_rate 1e-4 --num_train_epochs 20 --warmup_steps 1000\
         --gpu 2 --train_batch_size 4 --eval_batch_size 2 --max_seq_length 512 \
         --agg_strategy gmm,longnet --embed_dim 512 --vision_adaptor False --hierachical_token True --hierachical_adaptor True\
         --n_heads 32,16,8 --llm_requires_grad False --resume_from_checkpoint False \
@@ -67,7 +67,7 @@ accelerate launch --config_file=./accelerate_configs/deepspeed_zero2.yaml run_ws
 
 ## Train Step 2 ##
 ```
-accelerate launch --config_file=./accelerate_configs/deepspeed_zero2.yaml run_wsi.py --max_steps 20000 --warmup_steps 1000\
+accelerate launch --config_file=./accelerate_configs/deepspeed_zero2.yaml run_wsi.py --num_train_epochs 5 --warmup_steps 1000\
         --gpu 2 --train_batch_size 8 --eval_batch_size 2 --max_seq_length 256 \
         --agg_strategy gmm,longnet --embed_dim 512 --vision_adaptor False --hierachical_token True --hierachical_adaptor True\
         --n_heads 32,16,8 --llm_requires_grad True --resume_from_checkpoint False \