-
Notifications
You must be signed in to change notification settings - Fork 61
/
train_TopLandscape.sh
executable file
·63 lines (56 loc) · 2.19 KB
/
train_TopLandscape.sh
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
#!/bin/bash
set -x
source env.sh
echo "args: $@"
# set the dataset dir via `DATADIR_TopLandscape`
DATADIR=${DATADIR_TopLandscape}
[[ -z $DATADIR ]] && DATADIR='./datasets/TopLandscape'
# set a comment via `COMMENT`
suffix=${COMMENT}
# PN, PFN, PCNN, ParT
model=$1
extraopts=""
if [[ "$model" == "ParT" ]]; then
modelopts="networks/example_ParticleTransformer.py --use-amp --optimizer-option weight_decay 0.01"
lr="1e-3"
elif [[ "$model" == "ParT-FineTune" ]]; then
modelopts="networks/example_ParticleTransformer_finetune.py --use-amp --optimizer-option weight_decay 0.01"
lr="1e-4"
extraopts="--optimizer-option lr_mult (\"fc.*\",50) --lr-scheduler none --load-model-weights models/ParT_kin.pt"
elif [[ "$model" == "PN" ]]; then
modelopts="networks/example_ParticleNet.py"
lr="1e-2"
elif [[ "$model" == "PN-FineTune" ]]; then
modelopts="networks/example_ParticleNet_finetune.py"
lr="1e-3"
extraopts="--optimizer-option lr_mult (\"fc_out.*\",50) --lr-scheduler none --load-model-weights models/ParticleNet_kin.pt"
elif [[ "$model" == "PFN" ]]; then
modelopts="networks/example_PFN.py"
lr="2e-2"
extraopts="--batch-size 4096"
elif [[ "$model" == "PCNN" ]]; then
modelopts="networks/example_PCNN.py"
lr="2e-2"
extraopts="--batch-size 4096"
else
echo "Invalid model $model!"
exit 1
fi
# "kin"
FEATURE_TYPE=$2
[[ -z ${FEATURE_TYPE} ]] && FEATURE_TYPE="kin"
if [[ "${FEATURE_TYPE}" != "kin" ]]; then
echo "Invalid feature type ${FEATURE_TYPE}!"
exit 1
fi
weaver \
--data-train "${DATADIR}/train_file.parquet" \
--data-val "${DATADIR}/val_file.parquet" \
--data-test "${DATADIR}/test_file.parquet" \
--data-config data/TopLandscape/top_${FEATURE_TYPE}.yaml --network-config $modelopts \
--model-prefix training/TopLandscape/${model}/{auto}${suffix}/net \
--num-workers 1 --fetch-step 1 --in-memory \
--batch-size 512 --samples-per-epoch $((2400 * 512)) --samples-per-epoch-val $((800 * 512)) --num-epochs 20 --gpus 0 \
--start-lr $lr --optimizer ranger --log logs/TopLandscape_${model}_{auto}${suffix}.log --predict-output pred.root \
--tensorboard TopLandscape_${FEATURE_TYPE}_${model}${suffix} \
${extraopts} "${@:3}"