4D-DiT/scripts/train_S4_latent_sex.slurm 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
#!/bin/bash
#SBATCH -A m4727_g
#SBATCH -C gpu&hbm80g
#SBATCH --qos regular
#SBATCH --exclusive
#SBATCH -t 03:10:00
#SBATCH --nodes 1
#SBATCH --cpus-per-task=32
#SBATCH --ntasks-per-node=4
#SBATCH --gpus-per-node=4
#SBATCH --chdir='/global/homes/s/seojw/4D-DiT'
#SBATCH --output=/global/homes/s/seojw/4D-DiT/logs/output/R-%j-%x.out
#SBATCH -e /global/homes/s/seojw/4D-DiT/logs/error_log/VAE_%A.e
#SBATCH --mail-user=jungwoo.seo95@gmail.com

set +x

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

# --metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_full_train_set.csv' \
python train_latent.py \
    --distribution_type 'multi' \
    --model_dir /pscratch/sd/s/seojw/4D-DiT/checkpoints/latent-DiT \
    --experiment_name 'Latent_Swin4D-DiT-S4-sex_lr1e-4_real_full-VQGAN-309_59M' \
    --metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run1_LR.csv' \
    --latent_dif_model \
    --autoencoder_checkpoint '/pscratch/sd/s/seojw/4D-DiT/checkpoints/autoencoder/L1+percept_n=220836_depth=[6, 6]_embdim=216_lr=0.0005_ngpu=4_splitsave/AE_checkpoint_epoch_5.67.pth' \
    --bs 4 \
    --workers 8 \
    --lr 1e-4 \
    --depth 12 \
    --embed_dim 516 \
    --in_chans 432 \
    --num_head 12 \
    --saveIter 4 \
    --use_tb \
    --continue_training \

     #  --model_type 'DiT-S/4' \
     # --dataroot /data/HCP1200/fMRI/SwiFT_Prep/HCP_REST1_LR_MNI_to_TRs_z-norm_cropped \
     # --category sex \
     # --window_size 4 \
     # --window_block_indexes '0,3,6,9' \
     # --voxel_size 32 \


 #     --model_dir /data/connectome/seojw/4D-DiT/checkpoints \
 #     --metadata '/home/seojw/4D-DiT/HCP_sex_metadata_rest_run1_LR_debugdata.csv' \
# --metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run1_LR_debugdata_n300.csv' \
# --metadata '/global/homes/s/seojw/4D-DiT/HCP_sex_metadata_rest_run1_LR.csv' \