-
Notifications
You must be signed in to change notification settings - Fork 24
/
task_args.py
46 lines (37 loc) · 1.91 KB
/
task_args.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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import argparse
def forecasting_args(subparsers):
p = subparsers.add_parser('forecasting')
p.add_argument('--horizon', type=int, default=1, help='prediction horizon')
p.add_argument('--model_type', type=str, default='LSTM', choices=['LSTM', 'LSTNet', 'DSANet', 'AGCRN', 'TCN', 'TPA', 'Trans'])
p.add_argument('--dataset', type=str, default='solar')
# LSTM
p.add_argument('--lstm_n_layers', type=int, default=2)
p.add_argument('--lstm_hidden_size', type=int, default=64)
# DSANet
p.add_argument('--dsanet_local', type=int, default=3)
p.add_argument('--dsanet_n_kernels', type=int, default=32)
p.add_argument('--dsanet_w_kernel', type=int, default=1)
p.add_argument('--dsanet_d_model', type=int, default=512)
p.add_argument('--dsanet_d_inner', type=int, default=2048)
p.add_argument('--dsanet_d_k', type=int, default=64)
p.add_argument('--dsanet_d_v', type=int, default=64)
p.add_argument('--dsanet_n_head', type=int, default=8)
p.add_argument('--dsanet_n_layers', type=int, default=1)
p.add_argument('--dsanet_dropout', type=float, default=0.1)
# AGCRN
p.add_argument('--agcrn_n_layers', type=int, default=1)
p.add_argument('--agcrn_hidden_size', type=int, default=64)
p.add_argument('--agcrn_embed_dim', type=int, default=10)
# TCN
p.add_argument('--tcn_n_layers', type=int, default=9)
p.add_argument('--tcn_hidden_size', type=int, default=64)
p.add_argument('--tcn_dropout', type=int, default=0)
# TPA
p.add_argument('--tpa_n_layers', type=int, default=1)
p.add_argument('--tpa_hidden_size', type=int, default=64)
p.add_argument('--tpa_ar_len', type=int, default=24)
# Trans
p.add_argument('--trans_n_layers', type=int, default=3)
p.add_argument('--trans_n_head', type=int, default=8)
p.add_argument('--trans_hidden_size', type=int, default=256)
p.add_argument('--trans_kernel_size', type=int, default=6)