forked from ituvisionlab/climate-ai
-
Notifications
You must be signed in to change notification settings - Fork 0
/
args_parser.py
31 lines (25 loc) · 2.1 KB
/
args_parser.py
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
import argparse
def get_args():
# python cimp6_trainer_vars_reallyall.py --epochs 50 --template_path ../experiments/CMIP6-NestedUNet2all-0.0001 --model CMIP6-NestedUNet2all-0.0001 --batch_size 16 --device cuda:2 --years 3 --months 2 --weight_decay 0.0001 --lr 0.0001
parser = argparse.ArgumentParser()
parser.add_argument("--years", type=int, default=3, help="Number of lag years")
parser.add_argument("--months", type=int, default=2, help="Number of months that is going to be used for training. For historical: 6, 12, 18, 24, 30 36 For periodic: 1, 2")
parser.add_argument("--batch_size", type=int, default=2)
parser.add_argument("--epochs", type=int, default=40)
parser.add_argument("--device", type=str, default="cuda:0")
parser.add_argument("--prediction_month", type=int, default=1, help="Number of months to predict at output")
parser.add_argument("--positional_encoding", type=bool, default=False)
parser.add_argument("--template_path", type=str, default="../experiments/CMIP6-Dropout04_Calib", help="Path to save the model")
parser.add_argument("--model", type=str, default="CMIP6-Dropout04_Calib", help="Name of the model. Required for model_select function to chose the right model.")
parser.add_argument("--kl_coef", type=int, default=16, help="Only required when training bayesian models")
parser.add_argument("--climate_var", type=str , default="tas", choices=["pr", "psl", "tas", "uas", "vas", "zg500"])
parser.add_argument("--lr", type=float , default=1e-5)
parser.add_argument("--weight_decay", type=float , default=1e-5)
parser.add_argument("--dropout", type=bool , default=False)
parser.add_argument("--deeper", type=bool , default=False)
parser.add_argument("--ensembles", type=int , default=5)
parser.add_argument("--shell", "--transport", "--iopub", "-f", "--fff", "--ip", "--stdin", "--control", "--hb", "--Session.signature_scheme", "--Session.key", help="a dummy argument to fool ipython", default="1")
args = parser.parse_args()
if 0 < args.months < 3 and args.years > 4:
raise ValueError("Experiment not found.")
return args