4D-DiT/scripts/train_vqgan_4d_ddp_gan.sh at main · jungwoooseo/4D-DiT · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
#!/bin/bash
set +x

# correct this (same as #SBATCH --chdir)
cd /global/homes/s/seojw/4D-DiT
module load python
conda activate py39

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256
export NEPTUNE_API_TOKEN="eyJhcGlfYWRkcmVzcyI6Imh0dHBzOi8vYXBwLm5lcHR1bmUuYWkiLCJhcGlfdXJsIjoiaHR0cHM6Ly9hcHAubmVwdHVuZS5haSIsImFwaV9rZXkiOiIwZmQ5OGUxZi0xZDdhLTQ0ODQtYWY4Ni02NTJiYjE2YjdlZWIifQ=="

# --train_metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run1_LR_debugdata_n3000.csv' \
# --train_metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run1_LR.csv' \
# --train_metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_full_train_set.csv' \
# --train_metadata '/global/homes/s/seojw/4D-DiT/HCP_rest+task_metadata_full_train_set.csv' \

# --test_metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_full_validation_set_small.csv' \
# --test_metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run2_RL_testset.csv' \
# --test_metadata '/global/homes/s/seojw/4D-DiT/HCP_rest+task_metadata_full_validation_set_small.csv'

# python train_vqgan_4d_ddp.py \
#python train_vqgan_4d_ddp_evalsave.py \
#python train_vqgan_4d_ddp_gan_dev.py \

#python '/pscratch/sd/s/seojw/4D-DiT/checkpoints/autoencoder/L1+percept+inter_n=220836_depth=[4, 4, 4]_embdim=132_lr=0.0001_ngpu=4_splitsave/train_vqgan_4d_ddp_evalsave.py' \
#python train_vqgan_4d_ddp_evalsave_tmp20241216.py \
python train_vqgan_4d_ddp_gan_dev.py \
    --distribution_type 'multi' \
    --project_name jungwoo.seo/vqgan-4d \
    --train_metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run1_LR.csv' \
    --test_metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run2_RL_testset.csv' \
    --output_dir '/pscratch/sd/s/seojw/4D-DiT/checkpoints/autoencoder' \
    --num_workers 16 \
    --embed_dim 216 \
    --encoder_depth 6 6 \
    --decoder_depth 8 \
    --lr 1e-3 \
    --pixel_recon_loss L1 \
    --kl_weight 1e-6 \
    --recon_loss_weight 1.0 \
    --perceptual_weight 1.0 \
    --gan_weight_2d 0.0025 \
    --gan_weight_3d 0.0025 \
    --gan_feat_weight 0.1 \
    --disc_loss_type hinge \
    --discriminator_iter_start 50000 \
    --gradient_clipping \
    --accumulation_steps 1 \
    --bs 4 \
    --gan_continue_training \
    --gan_continue_training_checkpoint 4 \
    --gan_lr 5e-6
#--eval_interval 3 \
#    --gan_continue_training \
#    --gan_continue_training_checkpoint 2 \
#    --gan_lr 1e-3 \
#    --bs 4
#    --interframe_matching_brain_mask \