| import argparse |
| from s5.utils.util import str2bool |
| from s5.train import train |
| from s5.dataloading import Datasets |
|
|
| if __name__ == "__main__": |
|
|
| parser = argparse.ArgumentParser() |
|
|
| parser.add_argument("--USE_WANDB", type=str2bool, default=False, |
| help="log with wandb?") |
| parser.add_argument("--wandb_project", type=str, default=None, |
| help="wandb project name") |
| parser.add_argument("--wandb_entity", type=str, default=None, |
| help="wandb entity name, e.g. username") |
| parser.add_argument("--dir_name", type=str, default='./cache_dir', |
| help="name of directory where data is cached") |
| parser.add_argument("--dataset", type=str, choices=Datasets.keys(), |
| default='mnist-classification', |
| help="dataset name") |
|
|
| |
| parser.add_argument("--n_layers", type=int, default=6, |
| help="Number of layers in the network") |
| parser.add_argument("--d_model", type=int, default=128, |
| help="Number of features, i.e. H, " |
| "dimension of layer inputs/outputs") |
| parser.add_argument("--ssm_size_base", type=int, default=256, |
| help="SSM Latent size, i.e. P") |
| parser.add_argument("--blocks", type=int, default=8, |
| 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, |
| 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") |
|
|
| |
| 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=64, |
| help="batch size") |
| parser.add_argument("--epochs", type=int, default=100, |
| 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") |
|
|
| train(parser.parse_args()) |
|
|
