forked from peernagy/LOBS5
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathrun_train.py
159 lines (147 loc) · 8.29 KB
/
run_train.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
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
# CAVE: only for debugging purposes
#import os
#os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=48'
import argparse
from s5.utils.util import str2bool
from lob.train import train
from lob.dataloading import Datasets
#import tensorflow as tf
import os
import jax
import torch
import cProfile
if __name__ == "__main__":
#physical_devices = tf.config.list_physical_devices('GPU')
#tf.config.experimental.set_memory_growth(physical_devices[0], True)
#tf.config.experimental.set_visible_devices([], "GPU")
# no GPU use at all
#os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"]="false"
#os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".8"
torch.multiprocessing.set_start_method('spawn')
parser = argparse.ArgumentParser()
parser.add_argument("--USE_WANDB", type=str2bool, default=True,
help="log with wandb?")
parser.add_argument("--wandb_project", type=str, default="LOBS5",
help="wandb project name")
parser.add_argument("--wandb_entity", type=str, default="sasrey",
help="wandb entity name, e.g. username")
parser.add_argument("--dir_name", type=str, default='./data',
help="name of directory where data is cached")
parser.add_argument("--dataset", type=str, choices=Datasets.keys(),
default='lobster-prediction',
help="dataset name")
parser.add_argument("--masking", type=str, choices={'causal', 'random'},
default='causal', # random
help="causal or random masking of sequences")
parser.add_argument("--use_book_data", type=str2bool, default=False,
help="use book data in addition to message data")
parser.add_argument("--use_book_only", type=str2bool, default=False,
help="Use only book data and skip message data. Cannot be true with 'use_book_data.")
parser.add_argument("--use_simple_book", type=str2bool, default=False,
help="use raw price (-p0) and volume series instead of 'volume image representation'")
parser.add_argument("--book_transform", type=str2bool, default=False,
help="transform loaded book data to volume image repr. in dataloader")
parser.add_argument("--book_depth", type=int, default=500,
help="number of tick levels to use in book data [if book_transform=True]")
parser.add_argument("--restore", type=str,
help="if given restore from given checkpoint dir")
parser.add_argument("--restore_step", type=int)
parser.add_argument("--msg_seq_len", type=int, default=500, # 500
help="How many past messages to include in each sample")
parser.add_argument("--horizon_type", type=str, default='messages', choices=['messages','price-changes'],
help="Use n_messages or average number of price-changes to define horizon")
parser.add_argument("--prediction_horizon", type=int, default=10,
help="How many messages (events) or average price changes into the future to predict")
parser.add_argument("--zero_sequences", type=str2bool, default=False,
help="Whether or not to adjust sequences of data such that initial is 0. Offest-based.")
parser.add_argument("--n_data_workers", type=int, default=0,
help="number of workers used in DataLoader")
# Model Parameters
parser.add_argument("--n_message_layers", type=int, default=2, # 2
help="Number of layers after fusing message and book data")
parser.add_argument("--n_book_pre_layers", type=int, default=1, # 1
help="Number of layers taking in raw book data (before projecting dimensions)")
parser.add_argument("--n_book_post_layers", type=int, default=1, # 1
help="Number of book seq layers after projecting book data dimensions")
parser.add_argument("--n_layers", type=int, default=6, #6
help="Number of layers after fusing message and book data")
parser.add_argument("--d_model", type=int, default=32, #128, 32, 16
help="Number of features, i.e. H, "
"dimension of layer inputs/outputs")
parser.add_argument("--ssm_size_base", type=int, default=32, # 256
help="SSM Latent size, i.e. P")
parser.add_argument("--blocks", type=int, default=8, # 8, 4
help="How many blocks, J, to initialize with")
parser.add_argument("--C_init", type=str, default="trunc_standard_normal",
choices=["trunc_standard_normal", "lecun_normal", "complex_normal"],
help="Options for initialization of C: \\"
"trunc_standard_normal: sample from trunc. std. normal then multiply by V \\ " \
"lecun_normal sample from lecun normal, then multiply by V\\ " \
"complex_normal: sample directly from complex standard normal")
parser.add_argument("--discretization", type=str, default="zoh", choices=["zoh", "bilinear"])
parser.add_argument("--mode", type=str, default="pool", choices=["pool", "last"],
help="options: (for classification tasks) \\" \
" pool: mean pooling \\" \
"last: take last element")
parser.add_argument("--activation_fn", default="half_glu1", type=str,
choices=["full_glu", "half_glu1", "half_glu2", "gelu"])
parser.add_argument("--conj_sym", type=str2bool, default=True,
help="whether to enforce conjugate symmetry")
parser.add_argument("--clip_eigs", type=str2bool, default=False,
help="whether to enforce the left-half plane condition")
parser.add_argument("--bidirectional", type=str2bool, default=False, #False,
help="whether to use bidirectional model")
parser.add_argument("--dt_min", type=float, default=0.001,
help="min value to sample initial timescale params from")
parser.add_argument("--dt_max", type=float, default=0.1,
help="max value to sample initial timescale params from")
# Optimization Parameters
parser.add_argument("--prenorm", type=str2bool, default=True,
help="True: use prenorm, False: use postnorm")
parser.add_argument("--batchnorm", type=str2bool, default=True,
help="True: use batchnorm, False: use layernorm")
parser.add_argument("--bn_momentum", type=float, default=0.95,
help="batchnorm momentum")
parser.add_argument("--bsz", type=int, default=16, #64, (max 16 with full size)
help="batch size")
parser.add_argument("--num_devices", type=int, default=jax.device_count(),
help="number of devices (GPUs) to use")
parser.add_argument("--epochs", type=int, default=100, #100, 20
help="max number of epochs")
parser.add_argument("--early_stop_patience", type=int, default=1000,
help="number of epochs to continue training when val loss plateaus")
parser.add_argument("--ssm_lr_base", type=float, default=1e-3,
help="initial ssm learning rate")
parser.add_argument("--lr_factor", type=float, default=1,
help="global learning rate = lr_factor*ssm_lr_base")
parser.add_argument("--dt_global", type=str2bool, default=False,
help="Treat timescale parameter as global parameter or SSM parameter")
parser.add_argument("--lr_min", type=float, default=0,
help="minimum learning rate")
parser.add_argument("--cosine_anneal", type=str2bool, default=True,
help="whether to use cosine annealing schedule")
parser.add_argument("--warmup_end", type=int, default=1,
help="epoch to end linear warmup")
parser.add_argument("--lr_patience", type=int, default=1000000,
help="patience before decaying learning rate for lr_decay_on_val_plateau")
parser.add_argument("--reduce_factor", type=float, default=1.0,
help="factor to decay learning rate for lr_decay_on_val_plateau")
parser.add_argument("--p_dropout", type=float, default=0.0,
help="probability of dropout")
parser.add_argument("--weight_decay", type=float, default=0.05,
help="weight decay value")
parser.add_argument("--opt_config", type=str, default="standard", choices=['standard',
'BandCdecay',
'BfastandCdecay',
'noBCdecay'],
help="Opt configurations: \\ " \
"standard: no weight decay on B (ssm lr), weight decay on C (global lr) \\" \
"BandCdecay: weight decay on B (ssm lr), weight decay on C (global lr) \\" \
"BfastandCdecay: weight decay on B (global lr), weight decay on C (global lr) \\" \
"noBCdecay: no weight decay on B (ssm lr), no weight decay on C (ssm lr) \\")
parser.add_argument("--jax_seed", type=int, default=1919,
help="seed randomness")
#with jax.profiler.trace("/tmp/jax-trace", create_perfetto_link=True):
train(parser.parse_args())
#cProfile.run('train(parser.parse_args())')