online_main.py

# Copyright 2025 BrainX Ecosystem Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================


import os
import platform

os.environ['JAX_TRACEBACK_FILTERING'] = 'off'
from general_utils import MyArgumentParser
from dataset_shd import add_data_augment_args


def strtobool(val):
    val = val.lower()
    if val in ('y', 'yes', 't', 'true', 'on', '1'):
        return True
    elif val in ('n', 'no', 'f', 'false', 'off', '0'):
        return False
    else:
        raise ValueError(f"invalid truth value {val}")


num_worker = 0 if platform.system() == 'Windows' else 5


def add_training_options(parser):
    parser.add_argument("--load_exp_folder", type=str, default=None,
                        help="Path to experiment folder with a pretrained model to load. Note "
                             "that the same path will be used to store the current experiment.")
    parser.add_argument("--new_exp_folder", type=str, default=None, help="Path to output folder to store experiment.")
    parser.add_argument("--mode", type=str, default='train')
    parser.add_argument("--train_mode", type=str, default='batch', choices=['vmap', 'batch'])
    args, _ = parser.parse_known_args()

    parser.add_argument('--dataset', type=str, default='shd', choices=['shd', 'nmnist', 'gesture'])
    parser.add_argument('--data_length', type=int, default=100)
    parser.add_argument("--save_best", type=lambda x: bool(strtobool(str(x))), default=True,
                        help="If True, the model from the epoch with the highest validation "
                             "accuracy is saved, if False, no model is saved.")
    parser.add_argument("--batch_size", type=int, default=128, help="Number of input examples inside a single batch.")
    parser.add_argument("--nb_epochs", type=int, default=5, help="Number of training epochs.")
    parser.add_argument("--patience", type=int, default=20, help="early stopping patience.")
    parser.add_argument("--num_workers", type=int, default=num_worker)
    parser.add_argument("--seed", type=int, default=None,
                        help="Random seed for reproducibility (seeds brainstate.random and torch). "
                             "If None, behavior is unseeded (default).")
    parser.add_argument("--start_epoch", type=int, default=0,
                        help="Epoch number to start training at. Will be 0 if no pretrained "
                             "model is given. First epoch will be start_epoch+1.")
    parser.add_argument("--lr", type=float, default=1e-2,
                        help="Initial learning rate for training. The default value of 0.01 "
                             "is good for SHD and SC, but 0.001 seemed to work better for HD and SC.")
    parser.add_argument("--lr_step_size", type=int, default=40,
                        help="Number of epochs without progress before the learning rate gets decreased.")
    parser.add_argument("--lr_step_gamma", type=float, default=0.9,
                        help="Factor between 0 and 1 by which the learning rate gets "
                             "decreased when the scheduler patience is reached.")
    parser.add_argument("--use_augm", type=lambda x: bool(strtobool(str(x))), default=False,
                        help="Whether to use data augmentation or not.")
    add_data_augment_args(parser)
    return parser


def add_model_options(parser):
    parser.add_argument("--model_type", type=str, default="LIF", help="Type of ANN or SNN model.",
                        choices=["LIF", "adLIF", "RLIF", "RadLIF", "MLP", "RNN", "LiGRU", "GRU"])
    parser.add_argument("--surrogate", type=str, default="boxcar")
    parser.add_argument("--nb_layers", type=int, default=3, help="Number of layers (including readout layer).")
    parser.add_argument("--nb_hiddens", type=int, default=128, help="Number of neurons in all hidden layers.")
    parser.add_argument("--pdrop", type=float, default=0.1, help="Dropout rate, must be between 0 and 1.")
    parser.add_argument("--threshold", type=float, default=1.0)
    parser.add_argument("--train_threshold", type=float, default=0.95)
    parser.add_argument("--inp_scale", type=float, default=20.0)
    parser.add_argument("--rec_scale", type=float, default=5.0)
    parser.add_argument("--state_init", type=str, default='rand', choices=('rand', 'zero'),
                        help="Path to output folder to store experiment.")
    parser.add_argument("--normalization", type=str, default="none",
                        choices=["none", "batchnorm", "layernorm", 'dyt', 'rmsnorm', 'weightnorm'],
                        help="Type of normalization, every string different from batchnorm ")
    parser.add_argument("--use_bias", type=lambda x: bool(strtobool(str(x))), default=False,
                        help="Whether to include trainable bias with feedforward weights.")
    return parser


def parse_args():
    parser = MyArgumentParser(
        description="Model training on spiking speech commands datasets.",
        method='esd-rtrl'
    )
    parser = add_model_options(parser)
    parser = add_training_options(parser)
    args = parser.parse_args()
    return args


def main():
    """
    Runs model training/testing using the configuration specified
    by the parser arguments. Run `python online_main.py -h` for details.
    """

    # Get experiment configuration from parser
    args = parse_args()

    # Deterministic seeding for reproducible cross-version comparison.
    if getattr(args, 'seed', None) is not None:
        import brainstate
        brainstate.random.seed(args.seed)
        try:
            import torch
            torch.manual_seed(args.seed)
        except Exception:
            pass

    # Instantiate class for the desired experiment
    from online_model import Experiment
    experiment = Experiment(args)

    # Run experiment
    if args.mode == 'train':
        experiment.f_train()
    elif args.mode == 'test':
        experiment.f_test(8)
    else:
        raise ValueError("Mode must be either 'train' or 'test'")


if __name__ == "__main__":
    main()