SPDM/args.sh at master · watml/SPDM · 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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
# Default model paramters
BS=64

DATASET_NAME=$1
PRED=$2

NGPU=1

SIGMA_MAX=80.0
SIGMA_MIN=0.002
SIGMA_DATA=0.5
COV_XY=0
NUM_CH=256
NUM_HEADS=64
ATTN=32,16,8
SAMPLER=real-uniform
NUM_RES_BLOCKS=2
USE_16FP=True
ATTN_TYPE=reg
TEST_INTERVAL=5000
IN_CHANNELS=3
OUT_CHANNELS=3

# Arguments for group equivarient layers
G_EQUIV=False
G_OUTPUT="C4_S"
# Argument for group regularization
G_REG=False

# Arguments specific to each dataset
if [[ $DATASET_NAME == "e2h" ]]; then
    DATA_DIR=YOUR_DATASET_PATH
    DATASET=edges2handbags
    IMG_SIZE=64
    NUM_CH=192
    NUM_RES_BLOCKS=3
    EXP="e2h${IMG_SIZE}_${NUM_CH}d"
    SAVE_ITER=100000
elif [[ "$DATASET_NAME" == "lysto64" ]]; then
    DATA_DIR="/home/datasets/lysto64_random_crop_ddbm"
    DATASET="lysto64"
    IMG_SIZE=64
    NUM_CH=128
    NUM_RES_BLOCKS=3
    EXP="lysto64${IMG_SIZE}_${NUM_CH}d"
    SAVE_ITER=5000
elif [[ "$DATASET_NAME" == "ct_pet" ]]; then
    DATA_DIR="/data/ct_pet"
    DATASET="ct_pet"
    IMG_SIZE=256
    SIGMA_MAX=20.0
    SIGMA_MIN=0.0005
    INV_REG="H"
    EXP="ct_pet${IMG_SIZE}_${NUM_CH}d_INV_${INV_REG}"
    SAVE_ITER=20000
    AUTOENCODER_CKPT=/logs/pet_ct_ssim_lpips_l2_lr1e-5/checkpoint_700000.pth
fi

# Arguments for each of the possible diffusion noise schedules
if  [[ $PRED == "ve" ]]; then
    EXP+="_ve"
    COND=concat
elif  [[ $PRED == "vp" ]]; then
    EXP+="_vp"
    COND=concat
    BETA_D=2
    BETA_MIN=0.1
    SIGMA_MAX=1
    SIGMA_MIN=0.0001
elif  [[ $PRED == "ve_simple" ]]; then
    EXP+="_ve_simple"
    COND=concat
elif  [[ $PRED == "vp_simple" ]]; then
    EXP+="_vp_simple"
    COND=concat
    BETA_D=2
    BETA_MIN=0.1
    SIGMA_MAX=1
    SIGMA_MIN=0.0001
else
    echo "Not supported"
    exit 1
fi

# Arguments for image size in relation to batch size for training
if  [[ $IMG_SIZE == 256 ]]; then
    BS=16
elif  [[ $IMG_SIZE == 128 ]]; then
    BS=16
elif  [[ $IMG_SIZE == 128 ]]; then
    BS=14
elif  [[ $IMG_SIZE == 64 ]]; then
    BS=30
else
    echo "Not supported"
    exit 1
fi