diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/__init__.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/__init__.py index e541add3..e38488bc 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/common/__init__.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/__init__.py @@ -25,6 +25,7 @@ setup_training_environment, prepare_transforms, create_data_loaders, + shutdown_data_loaders, # Model creation and setup create_model, log_model_summary, diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py index ae82d88d..65e44154 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/common/train_base.py @@ -808,3 +808,89 @@ def create_data_loaders(dataset, dataset_test, train_sampler, test_sampler, args dataset_test, batch_size=args.batch_size, sampler=test_sampler, num_workers=args.workers, pin_memory=True if gpu > 0 else False, collate_fn=utils.collate_fn) return data_loader, data_loader_test + + +def _unregister_tracked_semaphores(*objects): + """Unregister multiprocessing semaphores from Python's resource_tracker. + + On Python <=3.11, ``_multiprocessing.SemLock``'s C dealloc calls + ``sem_close()`` but never ``resource_tracker.unregister()``. The + resource_tracker's ``atexit`` handler therefore reports every + semaphore ever created as "leaked". (Fixed in Python 3.12+ where + SemLock.__del__ calls unregister.) + + This function walks known multiprocessing container attributes + (Queue._rlock, Queue._wlock, Queue._sem, Event._cond, Event._flag, + etc.) to find underlying SemLock objects and unregisters their named + semaphores so the resource_tracker stays quiet. + """ + try: + from multiprocessing.resource_tracker import unregister + except ImportError: + return + + seen = set() + + def _scan(obj): + obj_id = id(obj) + if obj_id in seen: + return + seen.add(obj_id) + # Leaf: a Lock / Semaphore / BoundedSemaphore wrapping a SemLock + semlock = getattr(obj, '_semlock', None) + if semlock is not None: + name = getattr(semlock, 'name', None) + if name: + try: + unregister(name, "semaphore") + except Exception: + pass + return + # Recurse into known container attributes + for attr in ('_rlock', '_wlock', '_sem', '_lock', '_cond', '_flag'): + child = getattr(obj, attr, None) + if child is not None: + _scan(child) + + for obj in objects: + if obj is None: + continue + if isinstance(obj, (list, tuple)): + for item in obj: + _scan(item) + else: + _scan(obj) + + +def shutdown_data_loaders(*loaders): + """Explicitly shut down DataLoader worker processes to avoid leaked semaphore warnings. + + Must be called before exit when DataLoaders use num_workers > 0 (especially + on macOS where the 'spawn' start method tracks semaphores via resource_tracker). + Works for both persistent_workers=True and False. + + After joining workers, we explicitly unregister all POSIX named semaphores + owned by the iterator's multiprocessing Queues and Events from the + resource_tracker. On Python <=3.11, this unregister never happens + automatically (the C SemLock dealloc only calls sem_close, not + resource_tracker.unregister), so without this step the resource_tracker + warns about "leaked semaphore objects" at shutdown. + """ + import gc + for loader in loaders: + if not (hasattr(loader, '_iterator') and loader._iterator is not None): + continue + it = loader._iterator + try: + it._shutdown_workers() + except Exception: + pass + # Unregister all POSIX named semaphores from the resource_tracker + # so it does not report them as leaked at exit. + _unregister_tracked_semaphores( + getattr(it, '_index_queues', None), + getattr(it, '_data_queue', None), + getattr(it, '_workers_done_event', None), + ) + loader._iterator = None + gc.collect() diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/test_onnx.py b/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/test_onnx.py index 72b61848..deb0c46c 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/test_onnx.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/test_onnx.py @@ -35,6 +35,7 @@ from argparse import ArgumentParser from logging import getLogger +import numpy as np import onnxruntime as ort import pandas as pd import torch @@ -47,6 +48,7 @@ from tinyml_tinyverse.common.utils import misc_utils, utils, mdcl_utils from tinyml_tinyverse.common.utils.mdcl_utils import Logger from tinyml_tinyverse.common.utils.utils import get_confusion_matrix +from ..common.train_base import shutdown_data_loaders # Import common functions from base module from ..common.test_onnx_base import ( @@ -131,73 +133,76 @@ def main(gpu, args): data_loader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn) - - logger.info(f"Loading ONNX model: {args.model_path}") - ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) - - predicted = torch.tensor([]).to(device, non_blocking=True) - ground_truth = torch.tensor([]).to(device, non_blocking=True) - for batched_raw_data, batched_data, batched_target in data_loader: - batched_raw_data = batched_raw_data.to(device, non_blocking=True).long() - batched_data = batched_data.to(device, non_blocking=True).float() - batched_target = batched_target.to(device, non_blocking=True).long() - if transform: - batched_data = transform(batched_data) - if args.nn_for_feature_extraction: - for data in batched_raw_data: - predicted = torch.cat((predicted, torch.tensor( - ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy().astype(np.float32)})[0] - ).to(device))) - else: - for data in batched_data: - predicted = torch.cat((predicted, torch.tensor( - ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0] - ).to(device))) - ground_truth = torch.cat((ground_truth, batched_target)) - try: - mdcl_utils.create_dir(os.path.join(args.output_dir, 'post_training_analysis')) - logger.info("Plotting OvR Multiclass ROC score") - utils.plot_multiclass_roc(ground_truth, predicted, os.path.join(args.output_dir, 'post_training_analysis'), - label_map=dataset.inverse_label_map, phase='test') - logger.info("Plotting Class difference scores") - utils.plot_pairwise_differenced_class_scores(ground_truth, predicted, - os.path.join(args.output_dir, 'post_training_analysis'), - label_map=dataset.inverse_label_map, phase='test') - except Exception as e: - logger.warning(f"Post Training Analysis plots will not be generated because: {e}") - - metric = torcheval.metrics.MulticlassAccuracy() - # predicted = torch.argmax(predicted, dim=1) - metric.update(predicted, ground_truth) - logger = getLogger("root.main.test_data") - logger.info(f"Test Data Evaluation Accuracy: {metric.compute() * 100:.2f}%") - try: - logger.info( - f"Test Data Evaluation AUC ROC Score: {utils.get_au_roc(predicted, ground_truth, num_classes):.3f}") - except ValueError as e: - logger.warning("Not able to compute AUC ROC. Error: " + str(e)) - if len(torch.unique(ground_truth)) == 1: - logger.warning("Confusion Matrix can not be printed because only items of 1 class was present in test data") - else: - try: - confusion_matrix = get_confusion_matrix(predicted, ground_truth.type(torch.int64), - num_classes).cpu().numpy() - logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame( - confusion_matrix, columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], - index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), headers="keys", tablefmt='grid'))) - except ValueError as e: - logger.warning("Not able to compute Confusion Matrix. Error: " + str(e)) + + logger.info(f"Loading ONNX model: {args.model_path}") + ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) + + predicted = torch.tensor([]).to(device, non_blocking=True) + ground_truth = torch.tensor([]).to(device, non_blocking=True) + for batched_raw_data, batched_data, batched_target in data_loader: + batched_raw_data = batched_raw_data.to(device, non_blocking=True).long() + batched_data = batched_data.to(device, non_blocking=True).float() + batched_target = batched_target.to(device, non_blocking=True).long() + if transform: + batched_data = transform(batched_data) + if args.nn_for_feature_extraction: + for data in batched_raw_data: + predicted = torch.cat((predicted, torch.tensor( + ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy().astype(np.float32)})[0] + ).to(device))) + else: + for data in batched_data: + predicted = torch.cat((predicted, torch.tensor( + ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0] + ).to(device))) + ground_truth = torch.cat((ground_truth, batched_target)) try: - Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, - name="root.utils.print_file_level_classification_summary", append_log=True, console_log=False) - getLogger("root.utils.print_file_level_classification_summary").propagate = False - utils.print_file_level_classification_summary(dataset_test, predicted, ground_truth, "TestData") - logger.info(f"Generated File-level classification summary of test data in: {args.file_level_classification_log}") + mdcl_utils.create_dir(os.path.join(args.output_dir, 'post_training_analysis')) + logger.info("Plotting OvR Multiclass ROC score") + utils.plot_multiclass_roc(ground_truth, predicted, os.path.join(args.output_dir, 'post_training_analysis'), + label_map=dataset.inverse_label_map, phase='test') + logger.info("Plotting Class difference scores") + utils.plot_pairwise_differenced_class_scores(ground_truth, predicted, + os.path.join(args.output_dir, 'post_training_analysis'), + label_map=dataset.inverse_label_map, phase='test') except Exception as e: - logger.error(f"Failed to generate file-level classification summary: {str(e)}") + logger.warning(f"Post Training Analysis plots will not be generated because: {e}") + metric = torcheval.metrics.MulticlassAccuracy() + # predicted = torch.argmax(predicted, dim=1) + metric.update(predicted, ground_truth) + logger = getLogger("root.main.test_data") + logger.info(f"Test Data Evaluation Accuracy: {metric.compute() * 100:.2f}%") + try: + logger.info( + f"Test Data Evaluation AUC ROC Score: {utils.get_au_roc(predicted, ground_truth, num_classes):.3f}") + except ValueError as e: + logger.warning("Not able to compute AUC ROC. Error: " + str(e)) + if len(torch.unique(ground_truth)) == 1: + logger.warning("Confusion Matrix can not be printed because only items of 1 class was present in test data") + else: + try: + confusion_matrix = get_confusion_matrix(predicted, ground_truth.type(torch.int64), + num_classes).cpu().numpy() + logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame( + confusion_matrix, columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], + index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), headers="keys", tablefmt='grid'))) + except ValueError as e: + logger.warning("Not able to compute Confusion Matrix. Error: " + str(e)) + + try: + Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, + name="root.utils.print_file_level_classification_summary", append_log=True, console_log=False) + getLogger("root.utils.print_file_level_classification_summary").propagate = False + utils.print_file_level_classification_summary(dataset_test, predicted, ground_truth, "TestData") + logger.info(f"Generated File-level classification summary of test data in: {args.file_level_classification_log}") + except Exception as e: + logger.error(f"Failed to generate file-level classification summary: {str(e)}") + + finally: + shutdown_data_loaders(data_loader) return def run(args): diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py index 67047c0b..0d083836 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py @@ -113,6 +113,7 @@ apply_output_int_default, get_output_int_flag, load_onnx_for_inference, + shutdown_data_loaders, ) dataset_loader_dict = {'GenericImageDataset':GenericImageDataset} @@ -292,175 +293,179 @@ def main(gpu, args): logger.info("Loading data:") data_loader, data_loader_test = create_data_loaders(dataset, dataset_test, train_sampler, test_sampler, args, gpu) + try: - logger.info("Creating model") + logger.info("Creating model") - if args.load_saved_model == 'None': - if args.nas_enabled == 'True': - if args.quantization: - model = torch.load(os.path.join(os.path.dirname(args.output_dir), os.path.join('base', 'nas_model.pt')), weights_only=False) + if args.load_saved_model == 'None': + if args.nas_enabled == 'True': + if args.quantization: + model = torch.load(os.path.join(os.path.dirname(args.output_dir), os.path.join('base', 'nas_model.pt')), weights_only=False) + else: + nas_args = get_nas_args(args, data_loader, data_loader_test, num_classes, variables) + model = search_and_get_model(nas_args) + if not model: + logger.error("Please check on prior errors. NAS wasn't able to create a model") + sys.exit(1) + torch.save(model, os.path.join(args.output_dir, 'nas_model.pt')) else: - nas_args = get_nas_args(args, data_loader, data_loader_test, num_classes, variables) - model = search_and_get_model(nas_args) - if not model: - logger.error("Please check on prior errors. NAS wasn't able to create a model") - sys.exit(1) - torch.save(model, os.path.join(args.output_dir, 'nas_model.pt')) + model = models.get_model( + args.model, variables, num_classes, input_features=input_features, model_config=args.model_config, + model_spec=args.model_spec, + dual_op=args.dual_op) else: - model = models.get_model( - args.model, variables, num_classes, input_features=input_features, model_config=args.model_config, - model_spec=args.model_spec, - dual_op=args.dual_op) - else: - model = torch.load(args.load_saved_model, weights_only=False) - - if args.generic_model or args.nas_enabled: - summary_input_shape = (1,) + tuple(dataset.X.shape[1:]) - logger.info(f"Model summary input shape: {summary_input_shape}") - logger.info(f"{torchinfo.summary(model, summary_input_shape)}") - - model = load_pretrained_weights(model, args, logger) - - if handle_export_only(model, args, variables, input_features, logger): - return - - move_model_to_device(model, device, logger) - criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) - - model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) - optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) - resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) - - phase = 'QuantTrain' if args.quantization else 'FloatTrain' - logger.info("Start training") - start_time = timeit.default_timer() - best = dict(accuracy=0.0, f1=0, conf_matrix=dict(), epoch=None) - - # model = NeuralNetworkWithPreprocess - if args.nn_for_feature_extraction: - fe_model = models.FEModelLinear(dataset.X.shape[1], dataset.X_raw.shape[2], dataset.X.shape[2]).to(device) - fe_model = NeuralNetworkWithPreprocess(fe_model, None) - optimizer, lr_scheduler = setup_optimizer_and_scheduler(fe_model, args) - fe_model = utils.get_trained_feature_extraction_model( - fe_model, args, data_loader, data_loader_test, device, lr_scheduler, optimizer) - model = NeuralNetworkWithPreprocess(fe_model, model) - else: - model = NeuralNetworkWithPreprocess(None, model) - - # if output_int not set by user, then set it to default of task_type - if args.output_int == None: - args.output_int = True - - global _float_best_metric - sample_inputs = None - sample_targets = None - bsearch_float_metric = None - bsearch_example_inputs = None - if args.auto_quantization and args.quantization: - try: - sample_data_iter = iter(data_loader) - _, sample_data_fe, sample_targets_raw = next(sample_data_iter) - sample_inputs = sample_data_fe.float().to(device) - sample_targets = sample_targets_raw.long().to(device) - logger.info("Obtained sample data for auto quantization analysis") - except Exception as e: - logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") - bsearch_float_metric = _float_best_metric - try: - bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) - except Exception as e: - logger.warning(f"Could not get example inputs for binary search: {e}") + model = torch.load(args.load_saved_model, weights_only=False) - model = utils.quantization_wrapped_model( - model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, - args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, - calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, - eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, - task_type='classification', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, - autoquant_tolerance_classification=args.autoquant_tolerance_classification) + if args.generic_model or args.nas_enabled: + summary_input_shape = (1,) + tuple(dataset.X.shape[1:]) + logger.info(f"Model summary input shape: {summary_input_shape}") + logger.info(f"{torchinfo.summary(model, summary_input_shape)}") - for epoch in range(args.start_epoch, args.epochs): - if args.distributed: - train_sampler.set_epoch(epoch) + model = load_pretrained_weights(model, args, logger) - set_dataset_augmentation_enabled(dataset, True) + if handle_export_only(model, args, variables, input_features, logger): + return - utils.train_one_epoch_classification( - model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, - print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, - is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False, - nn_for_feature_extraction=args.nn_for_feature_extraction) + move_model_to_device(model, device, logger) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) - set_dataset_augmentation_enabled(dataset, False) - if not (args.quantization_method in ['PTQ'] and args.quantization): - lr_scheduler.step() - set_dataset_augmentation_enabled(dataset, False) - set_dataset_augmentation_enabled(dataset_test, False) - avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( - model, criterion, data_loader_test, device=device, transform=None, phase=phase, - num_classes=num_classes, dual_op=args.dual_op, nn_for_feature_extraction=args.nn_for_feature_extraction) - if model_ema: - avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( - model_ema, criterion, data_loader_test, device=device, transform=None, - log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op, - nn_for_feature_extraction=args.nn_for_feature_extraction) - if args.output_dir and avg_accuracy >= best['accuracy']: - logger.info(f"Epoch {epoch}: {avg_accuracy:.2f} (Val accuracy) >= {best['accuracy']:.2f} (So far best accuracy). Hence updating checkpoint.pth") - best['accuracy'], best['f1'], best['auc'], best['conf_matrix'], best['epoch'] = avg_accuracy, avg_f1, auc, avg_conf_matrix, epoch - best['predictions'], best['ground_truth'] = predictions, ground_truth - checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) - utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) - - if not args.quantization and args.auto_quantization: - _float_best_metric = best['accuracy'] / 100.0 - logger.info(f"Stored float best accuracy for binary search: {_float_best_metric:.4f}") - - # Log best epoch results - set_dataset_augmentation_enabled(dataset, False) - set_dataset_augmentation_enabled(dataset_test, False) - logger = getLogger(f"root.main.{phase}.BestEpoch") - logger.info("") - logger.info("Printing statistics of best epoch:") - logger.info(f"Best Epoch: {best['epoch']}") - logger.info(f"Acc@1 {best['accuracy']:.3f}") - logger.info(f"F1-Score {best['f1']:.3f}") - logger.info(f"AUC ROC Score {best['f1']:.3f}") - logger.info("") - logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame(best['conf_matrix'], - columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], - index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), - headers="keys", tablefmt='grid'))) - - Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, - name="root.utils.print_file_level_classification_summary", - append_log=True if args.quantization else False, console_log=False) - getLogger("root.utils.print_file_level_classification_summary").propagate = False - utils.print_file_level_classification_summary(dataset_test, best['predictions'], best['ground_truth'], phase) - logger.info(f"Generated file-level classification summary in: {args.file_level_classification_log}") - - # Export model - logger.info('Exporting model after training.') - if args.distributed is False or (args.distributed is True and int(os.environ['LOCAL_RANK']) == 0): + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) + optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) + resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) + + phase = 'QuantTrain' if args.quantization else 'FloatTrain' + logger.info("Start training") + start_time = timeit.default_timer() + best = dict(accuracy=0.0, f1=0, conf_matrix=dict(), epoch=None) + + # model = NeuralNetworkWithPreprocess if args.nn_for_feature_extraction: - example_input = next(iter(data_loader_test))[0] - input_shape = (1,) + dataset.X_raw.shape[1:] + fe_model = models.FEModelLinear(dataset.X.shape[1], dataset.X_raw.shape[2], dataset.X.shape[2]).to(device) + fe_model = NeuralNetworkWithPreprocess(fe_model, None) + optimizer, lr_scheduler = setup_optimizer_and_scheduler(fe_model, args) + fe_model = utils.get_trained_feature_extraction_model( + fe_model, args, data_loader, data_loader_test, device, lr_scheduler, optimizer) + model = NeuralNetworkWithPreprocess(fe_model, model) else: - example_input = next(iter(data_loader_test))[1] - input_shape = (1,) + dataset.X.shape[1:] - utils.export_model( - model, input_shape=input_shape, output_dir=args.output_dir, opset_version=args.opset_version, - quantization=args.quantization, example_input=example_input, generic_model=args.generic_model, - remove_hooks_for_jit=True if (args.quantization_method == TinyMLQuantizationMethod.PTQ and args.quantization) else False) - - log_training_time(start_time) - - if args.gen_golden_vectors: - + model = NeuralNetworkWithPreprocess(None, model) + + # if output_int not set by user, then set it to default of task_type + if args.output_int == None: + args.output_int = True + + global _float_best_metric + sample_inputs = None + sample_targets = None + bsearch_float_metric = None + bsearch_example_inputs = None + if args.auto_quantization and args.quantization: + try: + sample_data_iter = iter(data_loader) + _, sample_data_fe, sample_targets_raw = next(sample_data_iter) + sample_inputs = sample_data_fe.float().to(device) + sample_targets = sample_targets_raw.long().to(device) + logger.info("Obtained sample data for auto quantization analysis") + except Exception as e: + logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") + bsearch_float_metric = _float_best_metric + try: + bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) + except Exception as e: + logger.warning(f"Could not get example inputs for binary search: {e}") + + model = utils.quantization_wrapped_model( + model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, + args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, + calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, + eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, + task_type='classification', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, + autoquant_tolerance_classification=args.autoquant_tolerance_classification) + + for epoch in range(args.start_epoch, args.epochs): + if args.distributed: + train_sampler.set_epoch(epoch) + + set_dataset_augmentation_enabled(dataset, True) + + utils.train_one_epoch_classification( + model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, + print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, + is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False, + nn_for_feature_extraction=args.nn_for_feature_extraction) + + set_dataset_augmentation_enabled(dataset, False) + if not (args.quantization_method in ['PTQ'] and args.quantization): + lr_scheduler.step() + set_dataset_augmentation_enabled(dataset, False) + set_dataset_augmentation_enabled(dataset_test, False) + avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( + model, criterion, data_loader_test, device=device, transform=None, phase=phase, + num_classes=num_classes, dual_op=args.dual_op, nn_for_feature_extraction=args.nn_for_feature_extraction) + if model_ema: + avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( + model_ema, criterion, data_loader_test, device=device, transform=None, + log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op, + nn_for_feature_extraction=args.nn_for_feature_extraction) + if args.output_dir and avg_accuracy >= best['accuracy']: + logger.info(f"Epoch {epoch}: {avg_accuracy:.2f} (Val accuracy) >= {best['accuracy']:.2f} (So far best accuracy). Hence updating checkpoint.pth") + best['accuracy'], best['f1'], best['auc'], best['conf_matrix'], best['epoch'] = avg_accuracy, avg_f1, auc, avg_conf_matrix, epoch + best['predictions'], best['ground_truth'] = predictions, ground_truth + checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) + utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) + + if not args.quantization and args.auto_quantization: + _float_best_metric = best['accuracy'] / 100.0 + logger.info(f"Stored float best accuracy for binary search: {_float_best_metric:.4f}") + + # Log best epoch results set_dataset_augmentation_enabled(dataset, False) set_dataset_augmentation_enabled(dataset_test, False) - generate_golden_vector_dir(args.output_dir) - output_int = get_output_int_flag(args) - generate_golden_vectors(args.output_dir, dataset, output_int, args.generic_model, args.nn_for_feature_extraction) + logger = getLogger(f"root.main.{phase}.BestEpoch") + logger.info("") + logger.info("Printing statistics of best epoch:") + logger.info(f"Best Epoch: {best['epoch']}") + logger.info(f"Acc@1 {best['accuracy']:.3f}") + logger.info(f"F1-Score {best['f1']:.3f}") + logger.info(f"AUC ROC Score {best['f1']:.3f}") + logger.info("") + logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame(best['conf_matrix'], + columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], + index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), + headers="keys", tablefmt='grid'))) + + Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, + name="root.utils.print_file_level_classification_summary", + append_log=True if args.quantization else False, console_log=False) + getLogger("root.utils.print_file_level_classification_summary").propagate = False + utils.print_file_level_classification_summary(dataset_test, best['predictions'], best['ground_truth'], phase) + logger.info(f"Generated file-level classification summary in: {args.file_level_classification_log}") + + # Export model + logger.info('Exporting model after training.') + if args.distributed is False or (args.distributed is True and int(os.environ['LOCAL_RANK']) == 0): + if args.nn_for_feature_extraction: + example_input = next(iter(data_loader_test))[0] + input_shape = (1,) + dataset.X_raw.shape[1:] + else: + example_input = next(iter(data_loader_test))[1] + input_shape = (1,) + dataset.X.shape[1:] + utils.export_model( + model, input_shape=input_shape, output_dir=args.output_dir, opset_version=args.opset_version, + quantization=args.quantization, example_input=example_input, generic_model=args.generic_model, + remove_hooks_for_jit=True if (args.quantization_method == TinyMLQuantizationMethod.PTQ and args.quantization) else False) + + log_training_time(start_time) + + if args.gen_golden_vectors: + set_dataset_augmentation_enabled(dataset, False) + set_dataset_augmentation_enabled(dataset_test, False) + generate_golden_vector_dir(args.output_dir) + output_int = get_output_int_flag(args) + generate_golden_vectors(args.output_dir, dataset, output_int, args.generic_model, args.nn_for_feature_extraction) + + finally: + shutdown_data_loaders(data_loader, data_loader_test) + return def run(args): diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/test_onnx.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/test_onnx.py index ee7c729e..af68a950 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/test_onnx.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/test_onnx.py @@ -51,6 +51,7 @@ load_onnx_model, run_distributed_test, ) +from ..common.train_base import shutdown_data_loaders dataset_loader_dict = { 'GenericTSDataset': GenericTSDataset, @@ -120,24 +121,27 @@ def get_reconstruction_errors_stats(args): data_loader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=args.workers, pin_memory=True if args.gpu > 0 else False, collate_fn=utils.collate_fn) - - logger.info(f"Loading ONNX model: {args.model_path}") - ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) - - errors = torch.tensor([]).to(device, non_blocking=True) - for _, data, targets in data_loader: - data = data.to(device, non_blocking=True).float() - targets = targets.to(device, non_blocking=True).long() - batch_reconstruction_errors = torch.tensor([]).to(device, non_blocking=True) - for input, target_label in zip(data, targets): - input = input.unsqueeze(0).cpu().numpy() - output = torch.tensor(ort_sess.run([output_name], {input_name: input})[0]).to(device) - current_output_error = torch.mean((torch.from_numpy(input).to(device) - output) ** 2, dim=(1, 2, 3)) - batch_reconstruction_errors = torch.cat((batch_reconstruction_errors, current_output_error)) - errors = torch.cat((errors, batch_reconstruction_errors)) - - normal_error_mean = torch.mean(errors) - normal_error_std = torch.std(errors) + try: + + logger.info(f"Loading ONNX model: {args.model_path}") + ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) + + errors = torch.tensor([]).to(device, non_blocking=True) + for _, data, targets in data_loader: + data = data.to(device, non_blocking=True).float() + targets = targets.to(device, non_blocking=True).long() + batch_reconstruction_errors = torch.tensor([]).to(device, non_blocking=True) + for input, target_label in zip(data, targets): + input = input.unsqueeze(0).cpu().numpy() + output = torch.tensor(ort_sess.run([output_name], {input_name: input})[0]).to(device) + current_output_error = torch.mean((torch.from_numpy(input).to(device) - output) ** 2, dim=(1, 2, 3)) + batch_reconstruction_errors = torch.cat((batch_reconstruction_errors, current_output_error)) + errors = torch.cat((errors, batch_reconstruction_errors)) + + normal_error_mean = torch.mean(errors) + normal_error_std = torch.std(errors) + finally: + shutdown_data_loaders(data_loader) return normal_error_mean.cpu(), normal_error_std.cpu() @@ -168,108 +172,112 @@ def main(gpu, args): data_loader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=args.workers, pin_memory=True if gpu > 0 else False, collate_fn=utils.collate_fn) - - logger.info(f"Loading ONNX model: {args.model_path}") - ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) - - errors = torch.tensor([]).to(device, non_blocking=True) - ground_truth = torch.tensor([]).to(device, non_blocking=True) - - for _, data, targets in data_loader: - data = data.to(device, non_blocking=True).float() - targets = targets.to(device, non_blocking=True).long() - if transform: - data = transform(data) - batch_reconstruction_errors = torch.tensor([]).to(device, non_blocking=True) - batch_target_labels = torch.tensor([]).to(device, non_blocking=True) - for input, target_label in zip(data, targets): - input = input.unsqueeze(0).cpu().numpy() - output = torch.tensor(ort_sess.run([output_name], {input_name: input})[0]).to(device) - current_output_errors = torch.mean((torch.from_numpy(input).to(device) - output) ** 2, dim=(1, 2, 3)) - batch_reconstruction_errors = torch.cat((batch_reconstruction_errors, current_output_errors)) - batch_target_labels = torch.cat((batch_target_labels, target_label)) - errors = torch.cat((errors, batch_reconstruction_errors)) - ground_truth = torch.cat((ground_truth, batch_target_labels)) - - post_training_analysis_path = os.path.join(args.output_dir, 'post_training_analysis') - mdcl_utils.create_dir(post_training_analysis_path) - - # The classes folder in dataset should have two folders named Anomaly and Normal - anomaly_errors = errors[ground_truth == 0].cpu().numpy() - normal_errors = errors[ground_truth == 1].cpu().numpy() - logger.info("Plotting reconstructions errors") - - normal_train_mean, normal_train_std = get_reconstruction_errors_stats(args) - anomaly_test_mean = np.mean(anomaly_errors) - anomaly_test_std = np.std(anomaly_errors) - normal_test_mean = np.mean(normal_errors) - normal_test_std = np.std(normal_errors) - - # Results - logger.info(f"Reconstruction Error Statistics:") - logger.info(f"Normal training data - Mean: {normal_train_mean:.6f}, Std: {normal_train_std:.6f}") - logger.info(f"Anomaly test data - Mean: {anomaly_test_mean:.6f}, Std: {anomaly_test_std:.6f}") - logger.info(f"Normal test data - Mean: {normal_test_mean:.6f}, Std: {normal_test_std:.6f}") - - # Threshold - K is the number of standard deviations from the mean - all_k_values = [i * 0.5 for i in range(0, 10)] - results_data = [] - best_f1_score = 0 - best_f1_score_index = 0 - - for i, k in enumerate(all_k_values): - threshold = normal_train_mean + k * normal_train_std - results = get_model_performance(threshold, normal_errors, anomaly_errors) - results["k_value"] = k - results["threshold"] = float(threshold) - results_data.append(results) - if results["f1_score"] > best_f1_score: - best_f1_score_index = i - best_f1_score = results["f1_score"] - - csv_path = os.path.join(post_training_analysis_path, 'threshold_performance.csv') - with open(csv_path, 'w', newline='') as csvfile: - fieldnames = ['k_value', 'threshold', 'accuracy', 'precision', 'recall', - 'f1_score', 'false_positive_rate', 'true_positives', - 'true_negatives', 'false_positives', 'false_negatives'] - writer = csv.DictWriter(csvfile, fieldnames=fieldnames) - writer.writeheader() - for result in results_data: - for key in ['accuracy', 'precision', 'recall', 'f1_score', 'false_positive_rate']: - if key in result: - result[key] = round(result[key], 2) - writer.writerow(result) - - logger.info(f"Threshold performance data saved to {csv_path}") - - best_results = results_data[best_f1_score_index] - best_threshold = best_results["threshold"] - logger.info(f"Threshold for K = {best_results['k_value']} : {best_threshold:.6f}") - - utils.plot_reconstruction_errors(anomaly_errors, normal_errors, normal_train_mean, best_threshold, post_training_analysis_path) - utils.plot_reconstruction_errors(anomaly_errors, normal_errors, normal_train_mean, best_threshold, post_training_analysis_path, log_scale=True) - - logger.info(f"False positive rate: {best_results['false_positive_rate']:.2f}%") - logger.info(f"Anomaly detection rate (recall): {best_results['recall']:.2f}%") - logger.info(f"Accuracy: {best_results['accuracy']:.2f}%") - logger.info(f"Precision: {best_results['precision']:.2f}%") - logger.info(f"F1 Score: {best_results['f1_score']:.2f}%") + try: + + logger.info(f"Loading ONNX model: {args.model_path}") + ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) + + errors = torch.tensor([]).to(device, non_blocking=True) + ground_truth = torch.tensor([]).to(device, non_blocking=True) + + for _, data, targets in data_loader: + data = data.to(device, non_blocking=True).float() + targets = targets.to(device, non_blocking=True).long() + if transform: + data = transform(data) + batch_reconstruction_errors = torch.tensor([]).to(device, non_blocking=True) + batch_target_labels = torch.tensor([]).to(device, non_blocking=True) + for input, target_label in zip(data, targets): + input = input.unsqueeze(0).cpu().numpy() + output = torch.tensor(ort_sess.run([output_name], {input_name: input})[0]).to(device) + current_output_errors = torch.mean((torch.from_numpy(input).to(device) - output) ** 2, dim=(1, 2, 3)) + batch_reconstruction_errors = torch.cat((batch_reconstruction_errors, current_output_errors)) + batch_target_labels = torch.cat((batch_target_labels, target_label)) + errors = torch.cat((errors, batch_reconstruction_errors)) + ground_truth = torch.cat((ground_truth, batch_target_labels)) + + post_training_analysis_path = os.path.join(args.output_dir, 'post_training_analysis') + mdcl_utils.create_dir(post_training_analysis_path) + + # The classes folder in dataset should have two folders named Anomaly and Normal + anomaly_errors = errors[ground_truth == 0].cpu().numpy() + normal_errors = errors[ground_truth == 1].cpu().numpy() + logger.info("Plotting reconstructions errors") + + normal_train_mean, normal_train_std = get_reconstruction_errors_stats(args) + anomaly_test_mean = np.mean(anomaly_errors) + anomaly_test_std = np.std(anomaly_errors) + normal_test_mean = np.mean(normal_errors) + normal_test_std = np.std(normal_errors) + + # Results + logger.info(f"Reconstruction Error Statistics:") + logger.info(f"Normal training data - Mean: {normal_train_mean:.6f}, Std: {normal_train_std:.6f}") + logger.info(f"Anomaly test data - Mean: {anomaly_test_mean:.6f}, Std: {anomaly_test_std:.6f}") + logger.info(f"Normal test data - Mean: {normal_test_mean:.6f}, Std: {normal_test_std:.6f}") + + # Threshold - K is the number of standard deviations from the mean + all_k_values = [i * 0.5 for i in range(0, 10)] + results_data = [] + best_f1_score = 0 + best_f1_score_index = 0 + + for i, k in enumerate(all_k_values): + threshold = normal_train_mean + k * normal_train_std + results = get_model_performance(threshold, normal_errors, anomaly_errors) + results["k_value"] = k + results["threshold"] = float(threshold) + results_data.append(results) + if results["f1_score"] > best_f1_score: + best_f1_score_index = i + best_f1_score = results["f1_score"] + + csv_path = os.path.join(post_training_analysis_path, 'threshold_performance.csv') + with open(csv_path, 'w', newline='') as csvfile: + fieldnames = ['k_value', 'threshold', 'accuracy', 'precision', 'recall', + 'f1_score', 'false_positive_rate', 'true_positives', + 'true_negatives', 'false_positives', 'false_negatives'] + writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + writer.writeheader() + for result in results_data: + for key in ['accuracy', 'precision', 'recall', 'f1_score', 'false_positive_rate']: + if key in result: + result[key] = round(result[key], 2) + writer.writerow(result) + + logger.info(f"Threshold performance data saved to {csv_path}") + + best_results = results_data[best_f1_score_index] + best_threshold = best_results["threshold"] + logger.info(f"Threshold for K = {best_results['k_value']} : {best_threshold:.6f}") + + utils.plot_reconstruction_errors(anomaly_errors, normal_errors, normal_train_mean, best_threshold, post_training_analysis_path) + utils.plot_reconstruction_errors(anomaly_errors, normal_errors, normal_train_mean, best_threshold, post_training_analysis_path, log_scale=True) + + logger.info(f"False positive rate: {best_results['false_positive_rate']:.2f}%") + logger.info(f"Anomaly detection rate (recall): {best_results['recall']:.2f}%") + logger.info(f"Accuracy: {best_results['accuracy']:.2f}%") + logger.info(f"Precision: {best_results['precision']:.2f}%") + logger.info(f"F1 Score: {best_results['f1_score']:.2f}%") - # Create a 2x2 confusion matrix - confusion_matrix = np.array([ - [best_results['true_negatives'], best_results['false_positives']], - [best_results['false_negatives'], best_results['true_positives']] - ]) - - # Format using pandas and tabulate - confusion_matrix_df = pd.DataFrame( - confusion_matrix, - columns=["Predicted as: Normal", "Predicted as: Anomaly"], - index=["Ground Truth: Normal", "Ground Truth: Anomaly"] - ) - - logger.info('Confusion Matrix:\n {}'.format(tabulate(confusion_matrix_df, headers="keys", tablefmt='grid'))) + # Create a 2x2 confusion matrix + confusion_matrix = np.array([ + [best_results['true_negatives'], best_results['false_positives']], + [best_results['false_negatives'], best_results['true_positives']] + ]) + + # Format using pandas and tabulate + confusion_matrix_df = pd.DataFrame( + confusion_matrix, + columns=["Predicted as: Normal", "Predicted as: Anomaly"], + index=["Ground Truth: Normal", "Ground Truth: Anomaly"] + ) + + logger.info('Confusion Matrix:\n {}'.format(tabulate(confusion_matrix_df, headers="keys", tablefmt='grid'))) + + finally: + shutdown_data_loaders(data_loader) def run(args): diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/test_onnx_cls.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/test_onnx_cls.py index 1bbd04cb..66e55825 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/test_onnx_cls.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/test_onnx_cls.py @@ -35,6 +35,7 @@ from argparse import ArgumentParser from logging import getLogger +import numpy as np import onnxruntime as ort import pandas as pd import torch @@ -47,6 +48,7 @@ from tinyml_tinyverse.common.utils import misc_utils, utils, mdcl_utils from tinyml_tinyverse.common.utils.mdcl_utils import Logger from tinyml_tinyverse.common.utils.utils import get_confusion_matrix +from ..common.train_base import shutdown_data_loaders dataset_loader_dict = {'GenericTSDataset': GenericTSDataset} @@ -142,56 +144,59 @@ def main(gpu, args): dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn) - # data_loader_test = torch.utils.data.DataLoader( - # dataset_test, batch_size=args.batch_size, - # sampler=test_sampler, num_workers=args.workers, pin_memory=True, - # collate_fn=utils.collate_fn, ) - logger.info(f"Loading ONNX model: {args.model_path}") - if not args.generic_model: - utils.decrypt(args.model_path, utils.get_crypt_key()) - ort_sess = ort.InferenceSession(args.model_path) - if not args.generic_model: - utils.encrypt(args.model_path, utils.get_crypt_key()) - - input_name = ort_sess.get_inputs()[0].name - output_name = ort_sess.get_outputs()[0].name - predicted = torch.tensor([]).to(device, non_blocking=True) - ground_truth = torch.tensor([]).to(device, non_blocking=True) - for batched_raw_data, batched_data, batched_target in data_loader: - batched_raw_data = batched_raw_data.to(device, non_blocking=True).long() - batched_data = batched_data.to(device, non_blocking=True).float() - batched_target = batched_target.to(device, non_blocking=True).long() - if transform: - batched_data = transform(batched_data) - if args.nn_for_feature_extraction: - for data in batched_raw_data: - predicted = torch.cat((predicted, torch.tensor(ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy().astype(np.float32)})[0]).to(device))) + try: + # data_loader_test = torch.utils.data.DataLoader( + # dataset_test, batch_size=args.batch_size, + # sampler=test_sampler, num_workers=args.workers, pin_memory=True, + # collate_fn=utils.collate_fn, ) + logger.info(f"Loading ONNX model: {args.model_path}") + if not args.generic_model: + utils.decrypt(args.model_path, utils.get_crypt_key()) + ort_sess = ort.InferenceSession(args.model_path) + if not args.generic_model: + utils.encrypt(args.model_path, utils.get_crypt_key()) + + input_name = ort_sess.get_inputs()[0].name + output_name = ort_sess.get_outputs()[0].name + predicted = torch.tensor([]).to(device, non_blocking=True) + ground_truth = torch.tensor([]).to(device, non_blocking=True) + for batched_raw_data, batched_data, batched_target in data_loader: + batched_raw_data = batched_raw_data.to(device, non_blocking=True).long() + batched_data = batched_data.to(device, non_blocking=True).float() + batched_target = batched_target.to(device, non_blocking=True).long() + if transform: + batched_data = transform(batched_data) + if args.nn_for_feature_extraction: + for data in batched_raw_data: + predicted = torch.cat((predicted, torch.tensor(ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy().astype(np.float32)})[0]).to(device))) + else: + for data in batched_data: + predicted = torch.cat((predicted, torch.tensor(ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0]).to(device))) + ground_truth = torch.cat((ground_truth, batched_target)) + + mdcl_utils.create_dir(os.path.join(args.output_dir, 'post_training_analysis')) + # logger.info("Plotting OvR Multiclass ROC score") + # utils.plot_multiclass_roc(ground_truth.cpu().numpy(), predicted.cpu().numpy(), os.path.join(args.output_dir, 'post_training_analysis'), + # label_map=dataset.inverse_label_map, phase='test') + # logger.info("Plotting Class difference scores") + # utils.plot_pairwise_differenced_class_scores(ground_truth.cpu().numpy(), predicted.cpu().numpy(), os.path.join(args.output_dir, 'post_training_analysis'), + # label_map=dataset.inverse_label_map, phase='test') + metric = torcheval.metrics.MulticlassAccuracy() + metric.update(torch.argmax(predicted.squeeze(), dim=1), ground_truth.squeeze()) + logger = getLogger("root.main.test_data") + logger.info(f"Test Data Evaluation Accuracy: {metric.compute() * 100:.2f}%") + # logger.info( + # f"Test Data Evaluation AUC ROC Score: {utils.get_au_roc(predicted.type(torch.int64), ground_truth, num_classes):.3f}") + if len(torch.unique(ground_truth)) == 1: + logger.warning("Confusion Matrix can not be printed because only items of 1 class was present in test data") else: - for data in batched_data: - predicted = torch.cat((predicted, torch.tensor(ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0]).to(device))) - ground_truth = torch.cat((ground_truth, batched_target)) - - mdcl_utils.create_dir(os.path.join(args.output_dir, 'post_training_analysis')) - # logger.info("Plotting OvR Multiclass ROC score") - # utils.plot_multiclass_roc(ground_truth.cpu().numpy(), predicted.cpu().numpy(), os.path.join(args.output_dir, 'post_training_analysis'), - # label_map=dataset.inverse_label_map, phase='test') - # logger.info("Plotting Class difference scores") - # utils.plot_pairwise_differenced_class_scores(ground_truth.cpu().numpy(), predicted.cpu().numpy(), os.path.join(args.output_dir, 'post_training_analysis'), - # label_map=dataset.inverse_label_map, phase='test') - metric = torcheval.metrics.MulticlassAccuracy() - metric.update(torch.argmax(predicted.squeeze(), dim=1), ground_truth.squeeze()) - logger = getLogger("root.main.test_data") - logger.info(f"Test Data Evaluation Accuracy: {metric.compute() * 100:.2f}%") - # logger.info( - # f"Test Data Evaluation AUC ROC Score: {utils.get_au_roc(predicted.type(torch.int64), ground_truth, num_classes):.3f}") - if len(torch.unique(ground_truth)) == 1: - logger.warning("Confusion Matrix can not be printed because only items of 1 class was present in test data") - else: - confusion_matrix = get_confusion_matrix(predicted.squeeze().type(torch.int64), ground_truth.squeeze().type(torch.int64), - num_classes).cpu().numpy() - logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame( - confusion_matrix, columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], - index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), headers="keys", tablefmt='grid'))) + confusion_matrix = get_confusion_matrix(predicted.squeeze().type(torch.int64), ground_truth.squeeze().type(torch.int64), + num_classes).cpu().numpy() + logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame( + confusion_matrix, columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], + index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), headers="keys", tablefmt='grid'))) + finally: + shutdown_data_loaders(data_loader) return def run(args): diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/train.py index af399429..575f511d 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_anomalydetection/train.py @@ -77,6 +77,7 @@ get_output_int_flag, load_onnx_for_inference, create_data_loaders, + shutdown_data_loaders, ) dataset_loader_dict = {'GenericTSDatasetAD': GenericTSDatasetAD} @@ -203,123 +204,127 @@ def main(gpu, args): logger.info("Loading data:") data_loader, data_loader_test = create_data_loaders(dataset, dataset_test, train_sampler, test_sampler, args, gpu) - - logger.info("Creating model") - logger.info(f"Variables: {variables}, Input_features: {input_features}") - - # For anomaly detection, num_classes is input_features (autoencoder output) - model = models.get_model( - args.model, variables, num_classes=input_features, input_features=input_features, model_config=args.model_config, - model_spec=args.model_spec, - dual_op=args.dual_op) - - log_model_summary(model, args, variables, input_features, logger) - model = load_pretrained_weights(model, args, logger) - - # if output_int not set by user, then set it to default of task_type - if args.output_int == None: - args.output_int = False - - if handle_export_only(model, args, variables, input_features, logger): - return - - move_model_to_device(model, device, logger) - criterion = nn.MSELoss() - - global _float_best_metric - sample_inputs = None - sample_targets = None - bsearch_float_metric = None - bsearch_example_inputs = None - if args.auto_quantization and args.quantization: - try: - sample_data_iter = iter(data_loader) - _, sample_data_fe, sample_targets_raw = next(sample_data_iter) - sample_inputs = sample_data_fe.float().to(device) - sample_targets = sample_data_fe.float().to(device) - logger.info("Obtained sample data for auto quantization analysis") - except Exception as e: - logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") - bsearch_float_metric = _float_best_metric - try: - bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) - except Exception as e: - logger.warning(f"Could not get example inputs for binary search: {e}") - - model = utils.quantization_wrapped_model( - model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, - args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, - calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, - eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, - task_type='anomalydetection', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, - autoquant_tolerance_anomaly=args.autoquant_tolerance_anomaly) - - optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) - model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) - resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) - - phase = 'QuantTrain' if args.quantization else 'FloatTrain' - logger.info("Start training") - start_time = timeit.default_timer() - best = dict(mse=np.inf, epoch=None) - - for epoch in range(args.start_epoch, args.epochs): - if args.distributed: - train_sampler.set_epoch(epoch) - utils.train_one_epoch_anomalydetection( - model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, - print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, - is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False) - if not (args.quantization_method in ['PTQ'] and args.quantization): - lr_scheduler.step() - avg_mse = utils.evaluate_anomalydetection(model, criterion, data_loader_test, device=device, - transform=None, print_freq=args.print_freq, epoch=epoch, - phase=phase, num_classes=num_classes, dual_op=args.dual_op) - if model_ema: - avg_mse = utils.evaluate_anomalydetection( - model_ema, criterion, data_loader_test, device=device, transform=None, epoch=epoch, - log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op) - if args.output_dir and avg_mse <= best['mse']: - logger.info(f"Epoch[{epoch}]: {avg_mse:.6f} (Val MSE) <= {best['mse']:.6f} (So far least error). Hence updating checkpoint.pth") - best['mse'], best['epoch'] = avg_mse, epoch - checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) - utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) - - if not args.quantization and args.auto_quantization: - _float_best_metric = best['mse'] - logger.info(f"Stored float best MSE for binary search: {_float_best_metric:.4f}") - - # Log best epoch results - logger = getLogger(f"root.main.{phase}.BestEpoch") - logger.info("") - logger.info("Printing statistics of best epoch:") - logger.info(f"Best Epoch: {best['epoch']}") - logger.info(f"MSE {best['mse']:.3f}") - logger.info("") - - # Export model - logger.info('Exporting model after training.') - if args.distributed is False or (args.distributed is True and int(os.environ['LOCAL_RANK']) == 0): - utils.export_model( - model, input_shape=(1,) + dataset.X.shape[1:], output_dir=args.output_dir, opset_version=args.opset_version, - quantization=args.quantization, example_input=None, generic_model=args.generic_model, - remove_hooks_for_jit=True if (args.quantization_method == TinyMLQuantizationMethod.PTQ and args.quantization) else False) - if args.ondevice_training: - saved_onnx_path = os.path.join(args.output_dir, 'model.onnx') - ondevice_training.export_for_ondevice_training(saved_onnx_path, args) - ondevice_training.export_training_data(dataset, dataset_test, dataset_test_final, args) - - log_training_time(start_time) - - # Calculate threshold - model_path = os.path.join(args.output_dir, 'model.onnx') - error_mean, error_std = get_reconstruction_errors_stats(args.generic_model, model_path, args.device, data_loader) - threshold = error_mean + 3 * error_std - - if args.gen_golden_vectors: - generate_golden_vector_dir(args.output_dir) - output_int = get_output_int_flag(args) - generate_golden_vectors(args.output_dir, dataset, output_int, threshold, args.generic_model) + try: + + logger.info("Creating model") + logger.info(f"Variables: {variables}, Input_features: {input_features}") + + # For anomaly detection, num_classes is input_features (autoencoder output) + model = models.get_model( + args.model, variables, num_classes=input_features, input_features=input_features, model_config=args.model_config, + model_spec=args.model_spec, + dual_op=args.dual_op) + + log_model_summary(model, args, variables, input_features, logger) + model = load_pretrained_weights(model, args, logger) + + # if output_int not set by user, then set it to default of task_type + if args.output_int == None: + args.output_int = False + + if handle_export_only(model, args, variables, input_features, logger): + return + + move_model_to_device(model, device, logger) + criterion = nn.MSELoss() + + global _float_best_metric + sample_inputs = None + sample_targets = None + bsearch_float_metric = None + bsearch_example_inputs = None + if args.auto_quantization and args.quantization: + try: + sample_data_iter = iter(data_loader) + _, sample_data_fe, sample_targets_raw = next(sample_data_iter) + sample_inputs = sample_data_fe.float().to(device) + sample_targets = sample_data_fe.float().to(device) + logger.info("Obtained sample data for auto quantization analysis") + except Exception as e: + logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") + bsearch_float_metric = _float_best_metric + try: + bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) + except Exception as e: + logger.warning(f"Could not get example inputs for binary search: {e}") + + model = utils.quantization_wrapped_model( + model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, + args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, + calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, + eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, + task_type='anomalydetection', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, + autoquant_tolerance_anomaly=args.autoquant_tolerance_anomaly) + + optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) + resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) + + phase = 'QuantTrain' if args.quantization else 'FloatTrain' + logger.info("Start training") + start_time = timeit.default_timer() + best = dict(mse=np.inf, epoch=None) + + for epoch in range(args.start_epoch, args.epochs): + if args.distributed: + train_sampler.set_epoch(epoch) + utils.train_one_epoch_anomalydetection( + model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, + print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, + is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False) + if not (args.quantization_method in ['PTQ'] and args.quantization): + lr_scheduler.step() + avg_mse = utils.evaluate_anomalydetection(model, criterion, data_loader_test, device=device, + transform=None, print_freq=args.print_freq, epoch=epoch, + phase=phase, num_classes=num_classes, dual_op=args.dual_op) + if model_ema: + avg_mse = utils.evaluate_anomalydetection( + model_ema, criterion, data_loader_test, device=device, transform=None, epoch=epoch, + log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op) + if args.output_dir and avg_mse <= best['mse']: + logger.info(f"Epoch[{epoch}]: {avg_mse:.6f} (Val MSE) <= {best['mse']:.6f} (So far least error). Hence updating checkpoint.pth") + best['mse'], best['epoch'] = avg_mse, epoch + checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) + utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) + + if not args.quantization and args.auto_quantization: + _float_best_metric = best['mse'] + logger.info(f"Stored float best MSE for binary search: {_float_best_metric:.4f}") + + # Log best epoch results + logger = getLogger(f"root.main.{phase}.BestEpoch") + logger.info("") + logger.info("Printing statistics of best epoch:") + logger.info(f"Best Epoch: {best['epoch']}") + logger.info(f"MSE {best['mse']:.3f}") + logger.info("") + + # Export model + logger.info('Exporting model after training.') + if args.distributed is False or (args.distributed is True and int(os.environ['LOCAL_RANK']) == 0): + utils.export_model( + model, input_shape=(1,) + dataset.X.shape[1:], output_dir=args.output_dir, opset_version=args.opset_version, + quantization=args.quantization, example_input=None, generic_model=args.generic_model, + remove_hooks_for_jit=True if (args.quantization_method == TinyMLQuantizationMethod.PTQ and args.quantization) else False) + if args.ondevice_training: + saved_onnx_path = os.path.join(args.output_dir, 'model.onnx') + ondevice_training.export_for_ondevice_training(saved_onnx_path, args) + ondevice_training.export_training_data(dataset, dataset_test, dataset_test_final, args) + + log_training_time(start_time) + + # Calculate threshold + model_path = os.path.join(args.output_dir, 'model.onnx') + error_mean, error_std = get_reconstruction_errors_stats(args.generic_model, model_path, args.device, data_loader) + threshold = error_mean + 3 * error_std + + if args.gen_golden_vectors: + generate_golden_vector_dir(args.output_dir) + output_int = get_output_int_flag(args) + generate_golden_vectors(args.output_dir, dataset, output_int, threshold, args.generic_model) + + finally: + shutdown_data_loaders(data_loader, data_loader_test) def run(args): diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/test_onnx.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/test_onnx.py index fb0d8aad..5fa659ea 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/test_onnx.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/test_onnx.py @@ -53,6 +53,7 @@ load_onnx_model, run_distributed_test, ) +from ..common.train_base import shutdown_data_loaders dataset_loader_dict = {'GenericTSDataset': GenericTSDataset} @@ -105,73 +106,76 @@ def main(gpu, args): data_loader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn) - - logger.info(f"Loading ONNX model: {args.model_path}") - ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) - - predicted = torch.tensor([]).to(device, non_blocking=True) - ground_truth = torch.tensor([]).to(device, non_blocking=True) - - for batched_raw_data, batched_data, batched_target in data_loader: - batched_raw_data = batched_raw_data.to(device, non_blocking=True).long() - batched_data = batched_data.to(device, non_blocking=True).float() - batched_target = batched_target.to(device, non_blocking=True).long() - if transform: - batched_data = transform(batched_data) - if args.nn_for_feature_extraction: - for data in batched_raw_data: - predicted = torch.cat((predicted, torch.tensor( - ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy().astype(np.float32)})[0] - ).to(device))) - else: - for data in batched_data: - predicted = torch.cat((predicted, torch.tensor( - ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0] - ).to(device))) - ground_truth = torch.cat((ground_truth, batched_target)) - try: - mdcl_utils.create_dir(os.path.join(args.output_dir, 'post_training_analysis')) - logger.info("Plotting OvR Multiclass ROC score") - utils.plot_multiclass_roc(ground_truth, predicted, os.path.join(args.output_dir, 'post_training_analysis'), - label_map=dataset.inverse_label_map, phase='test') - logger.info("Plotting Class difference scores") - utils.plot_pairwise_differenced_class_scores(ground_truth, predicted, - os.path.join(args.output_dir, 'post_training_analysis'), - label_map=dataset.inverse_label_map, phase='test') - except Exception as e: - logger.warning(f"Post Training Analysis plots will not be generated because: {e}") - - metric = torcheval.metrics.MulticlassAccuracy() - metric.update(predicted, ground_truth) - logger = getLogger("root.main.test_data") - logger.info(f"Test Data Evaluation Accuracy: {metric.compute() * 100:.2f}%") - try: - logger.info(f"Test Data Evaluation AUC ROC Score: {utils.get_au_roc(predicted, ground_truth, num_classes):.3f}") - except ValueError as e: - logger.warning("Not able to compute AUC ROC. Error: " + str(e)) + logger.info(f"Loading ONNX model: {args.model_path}") + ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) + + predicted = torch.tensor([]).to(device, non_blocking=True) + ground_truth = torch.tensor([]).to(device, non_blocking=True) + + for batched_raw_data, batched_data, batched_target in data_loader: + batched_raw_data = batched_raw_data.to(device, non_blocking=True).long() + batched_data = batched_data.to(device, non_blocking=True).float() + batched_target = batched_target.to(device, non_blocking=True).long() + if transform: + batched_data = transform(batched_data) + if args.nn_for_feature_extraction: + for data in batched_raw_data: + predicted = torch.cat((predicted, torch.tensor( + ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy().astype(np.float32)})[0] + ).to(device))) + else: + for data in batched_data: + predicted = torch.cat((predicted, torch.tensor( + ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0] + ).to(device))) + ground_truth = torch.cat((ground_truth, batched_target)) - if len(torch.unique(ground_truth)) == 1: - logger.warning("Confusion Matrix can not be printed because only items of 1 class was present in test data") - else: try: - confusion_matrix = get_confusion_matrix(predicted, ground_truth.type(torch.int64), num_classes).cpu().numpy() - logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame( - confusion_matrix, columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], - index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), headers="keys", tablefmt='grid'))) - except ValueError as e: - logger.warning("Not able to compute Confusion Matrix. Error: " + str(e)) + mdcl_utils.create_dir(os.path.join(args.output_dir, 'post_training_analysis')) + logger.info("Plotting OvR Multiclass ROC score") + utils.plot_multiclass_roc(ground_truth, predicted, os.path.join(args.output_dir, 'post_training_analysis'), + label_map=dataset.inverse_label_map, phase='test') + logger.info("Plotting Class difference scores") + utils.plot_pairwise_differenced_class_scores(ground_truth, predicted, + os.path.join(args.output_dir, 'post_training_analysis'), + label_map=dataset.inverse_label_map, phase='test') + except Exception as e: + logger.warning(f"Post Training Analysis plots will not be generated because: {e}") + + metric = torcheval.metrics.MulticlassAccuracy() + metric.update(predicted, ground_truth) + logger = getLogger("root.main.test_data") + logger.info(f"Test Data Evaluation Accuracy: {metric.compute() * 100:.2f}%") try: - Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, - name="root.utils.print_file_level_classification_summary", append_log=True, console_log=False) - getLogger("root.utils.print_file_level_classification_summary").propagate = False - utils.print_file_level_classification_summary(dataset_test, predicted, ground_truth, "TestData") - logger.info(f"Generated File-level classification summary of test data in: {args.file_level_classification_log}") - except Exception as e: - logger.error(f"Failed to generate file-level classification summary: {str(e)}") + logger.info(f"Test Data Evaluation AUC ROC Score: {utils.get_au_roc(predicted, ground_truth, num_classes):.3f}") + except ValueError as e: + logger.warning("Not able to compute AUC ROC. Error: " + str(e)) + if len(torch.unique(ground_truth)) == 1: + logger.warning("Confusion Matrix can not be printed because only items of 1 class was present in test data") + else: + try: + confusion_matrix = get_confusion_matrix(predicted, ground_truth.type(torch.int64), num_classes).cpu().numpy() + logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame( + confusion_matrix, columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], + index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), headers="keys", tablefmt='grid'))) + except ValueError as e: + logger.warning("Not able to compute Confusion Matrix. Error: " + str(e)) + + try: + Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, + name="root.utils.print_file_level_classification_summary", append_log=True, console_log=False) + getLogger("root.utils.print_file_level_classification_summary").propagate = False + utils.print_file_level_classification_summary(dataset_test, predicted, ground_truth, "TestData") + logger.info(f"Generated File-level classification summary of test data in: {args.file_level_classification_log}") + except Exception as e: + logger.error(f"Failed to generate file-level classification summary: {str(e)}") + + finally: + shutdown_data_loaders(data_loader) return diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py index 3a858474..276f3264 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py @@ -73,6 +73,7 @@ setup_training_environment, prepare_transforms, create_data_loaders, + shutdown_data_loaders, log_model_summary, load_pretrained_weights, setup_optimizer_and_scheduler, @@ -237,160 +238,164 @@ def main(gpu, args): logger.info("Loading data:") data_loader, data_loader_test = create_data_loaders(dataset, dataset_test, train_sampler, test_sampler, args, gpu) + try: - logger.info("Creating model") - if args.load_saved_model == 'None': - if args.nas_enabled == 'True': - if args.quantization: - model = torch.load(os.path.join(os.path.dirname(args.output_dir), os.path.join('base', 'nas_model.pt')), weights_only=False) + logger.info("Creating model") + if args.load_saved_model == 'None': + if args.nas_enabled == 'True': + if args.quantization: + model = torch.load(os.path.join(os.path.dirname(args.output_dir), os.path.join('base', 'nas_model.pt')), weights_only=False) + else: + nas_args = get_nas_args(args, data_loader, data_loader_test, num_classes, variables) + model = search_and_get_model(nas_args) + if not model: + logger.error("Please check on prior errors. NAS wasn't able to create a model") + sys.exit(1) + torch.save(model, os.path.join(args.output_dir, 'nas_model.pt')) else: - nas_args = get_nas_args(args, data_loader, data_loader_test, num_classes, variables) - model = search_and_get_model(nas_args) - if not model: - logger.error("Please check on prior errors. NAS wasn't able to create a model") - sys.exit(1) - torch.save(model, os.path.join(args.output_dir, 'nas_model.pt')) + model = models.get_model( + args.model, variables, num_classes, input_features=input_features, model_config=args.model_config, + model_spec=args.model_spec, + dual_op=args.dual_op) else: - model = models.get_model( - args.model, variables, num_classes, input_features=input_features, model_config=args.model_config, - model_spec=args.model_spec, - dual_op=args.dual_op) - else: - model = torch.load(args.load_saved_model, weights_only=False) + model = torch.load(args.load_saved_model, weights_only=False) - if args.generic_model or args.nas_enabled: - log_model_summary(model, args, variables, input_features, logger) + if args.generic_model or args.nas_enabled: + log_model_summary(model, args, variables, input_features, logger) - model = load_pretrained_weights(model, args, logger) + model = load_pretrained_weights(model, args, logger) - if handle_export_only(model, args, variables, input_features, logger): - return + if handle_export_only(model, args, variables, input_features, logger): + return - move_model_to_device(model, device, logger) - criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + move_model_to_device(model, device, logger) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) - model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) - optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) - resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) + optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) + resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) - phase = 'QuantTrain' if args.quantization else 'FloatTrain' - logger.info("Start training") - start_time = timeit.default_timer() - best = dict(accuracy=0.0, f1=0, conf_matrix=dict(), epoch=None) + phase = 'QuantTrain' if args.quantization else 'FloatTrain' + logger.info("Start training") + start_time = timeit.default_timer() + best = dict(accuracy=0.0, f1=0, conf_matrix=dict(), epoch=None) - # Handle nn_for_feature_extraction - if args.nn_for_feature_extraction: - fe_model = models.FEModelLinear(dataset.X.shape[1], dataset.X_raw.shape[2], dataset.X.shape[2]).to(device) - fe_model = NeuralNetworkWithPreprocess(fe_model, None) - optimizer, lr_scheduler = setup_optimizer_and_scheduler(fe_model, args) - fe_model = utils.get_trained_feature_extraction_model( - fe_model, args, data_loader, data_loader_test, device, lr_scheduler, optimizer) - model = NeuralNetworkWithPreprocess(fe_model, model) - else: - model = NeuralNetworkWithPreprocess(None, model) - - # if output_int not set by user, then set it to default of task_type - if args.output_int == None: - args.output_int = True - - global _float_best_metric - sample_inputs = None - sample_targets = None - bsearch_float_metric = None - bsearch_example_inputs = None - if args.auto_quantization and args.quantization: - try: - sample_data_iter = iter(data_loader) - _, sample_data_fe, sample_targets_raw = next(sample_data_iter) - sample_inputs = sample_data_fe.float().to(device) - sample_targets = sample_targets_raw.long().to(device) - logger.info("Obtained sample data for auto quantization analysis") - except Exception as e: - logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") - bsearch_float_metric = _float_best_metric - try: - bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) - except Exception as e: - logger.warning(f"Could not get example inputs for binary search: {e}") - - model = utils.quantization_wrapped_model( - model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, - args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, - calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, - eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, - task_type='classification', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, - autoquant_tolerance_classification=args.autoquant_tolerance_classification) - - for epoch in range(args.start_epoch, args.epochs): - if args.distributed: - train_sampler.set_epoch(epoch) - utils.train_one_epoch_classification( - model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, - print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, - is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False, - nn_for_feature_extraction=args.nn_for_feature_extraction) - if not (args.quantization_method in ['PTQ'] and args.quantization): - lr_scheduler.step() - avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( - model, criterion, data_loader_test, device=device, transform=None, phase=phase, - num_classes=num_classes, dual_op=args.dual_op, nn_for_feature_extraction=args.nn_for_feature_extraction) - if model_ema: - avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( - model_ema, criterion, data_loader_test, device=device, transform=None, - log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op, - nn_for_feature_extraction=args.nn_for_feature_extraction) - if args.output_dir and avg_accuracy >= best['accuracy']: - logger.info(f"Epoch {epoch}: {avg_accuracy:.2f} (Val accuracy) >= {best['accuracy']:.2f} (So far best accuracy). Hence updating checkpoint.pth") - best['accuracy'], best['f1'], best['auc'], best['conf_matrix'], best['epoch'] = avg_accuracy, avg_f1, auc, avg_conf_matrix, epoch - best['predictions'], best['ground_truth'] = predictions, ground_truth - checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) - utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) - - if not args.quantization and args.auto_quantization: - _float_best_metric = best['accuracy'] / 100.0 - logger.info(f"Stored float best accuracy for binary search: {_float_best_metric:.4f}") - - # Log best epoch results - logger = getLogger(f"root.main.{phase}.BestEpoch") - logger.info("") - logger.info("Printing statistics of best epoch:") - logger.info(f"Best Epoch: {best['epoch']}") - logger.info(f"Acc@1 {best['accuracy']:.3f}") - logger.info(f"F1-Score {best['f1']:.3f}") - logger.info(f"AUC ROC Score {best['f1']:.3f}") - logger.info("") - logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame(best['conf_matrix'], - columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], - index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), - headers="keys", tablefmt='grid'))) - - Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, - name="root.utils.print_file_level_classification_summary", - append_log=True if args.quantization else False, console_log=False) - getLogger("root.utils.print_file_level_classification_summary").propagate = False - utils.print_file_level_classification_summary(dataset_test, best['predictions'], best['ground_truth'], phase) - logger.info(f"Generated file-level classification summary in: {args.file_level_classification_log}") - - # Export model - logger.info('Exporting model after training.') - if args.distributed is False or (args.distributed is True and int(os.environ['LOCAL_RANK']) == 0): + # Handle nn_for_feature_extraction if args.nn_for_feature_extraction: - example_input = next(iter(data_loader_test))[0] - input_shape = (1,) + dataset.X_raw.shape[1:] + fe_model = models.FEModelLinear(dataset.X.shape[1], dataset.X_raw.shape[2], dataset.X.shape[2]).to(device) + fe_model = NeuralNetworkWithPreprocess(fe_model, None) + optimizer, lr_scheduler = setup_optimizer_and_scheduler(fe_model, args) + fe_model = utils.get_trained_feature_extraction_model( + fe_model, args, data_loader, data_loader_test, device, lr_scheduler, optimizer) + model = NeuralNetworkWithPreprocess(fe_model, model) else: - example_input = next(iter(data_loader_test))[1] - input_shape = (1,) + dataset.X.shape[1:] - utils.export_model( - model, input_shape=input_shape, output_dir=args.output_dir, opset_version=args.opset_version, - quantization=args.quantization, example_input=example_input, generic_model=args.generic_model, - remove_hooks_for_jit=True if (args.quantization_method == TinyMLQuantizationMethod.PTQ and args.quantization) else False) - - log_training_time(start_time) - - if args.gen_golden_vectors: - generate_golden_vector_dir(args.output_dir) - output_int = get_output_int_flag(args) - generate_golden_vectors(args.output_dir, dataset, output_int, args.generic_model, args.nn_for_feature_extraction) + model = NeuralNetworkWithPreprocess(None, model) + + # if output_int not set by user, then set it to default of task_type + if args.output_int == None: + args.output_int = True + + global _float_best_metric + sample_inputs = None + sample_targets = None + bsearch_float_metric = None + bsearch_example_inputs = None + if args.auto_quantization and args.quantization: + try: + sample_data_iter = iter(data_loader) + _, sample_data_fe, sample_targets_raw = next(sample_data_iter) + sample_inputs = sample_data_fe.float().to(device) + sample_targets = sample_targets_raw.long().to(device) + logger.info("Obtained sample data for auto quantization analysis") + except Exception as e: + logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") + bsearch_float_metric = _float_best_metric + try: + bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) + except Exception as e: + logger.warning(f"Could not get example inputs for binary search: {e}") + + model = utils.quantization_wrapped_model( + model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, + args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, + calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, + eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, + task_type='classification', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, + autoquant_tolerance_classification=args.autoquant_tolerance_classification) + + for epoch in range(args.start_epoch, args.epochs): + if args.distributed: + train_sampler.set_epoch(epoch) + utils.train_one_epoch_classification( + model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, + print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, + is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False, + nn_for_feature_extraction=args.nn_for_feature_extraction) + if not (args.quantization_method in ['PTQ'] and args.quantization): + lr_scheduler.step() + avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( + model, criterion, data_loader_test, device=device, transform=None, phase=phase, + num_classes=num_classes, dual_op=args.dual_op, nn_for_feature_extraction=args.nn_for_feature_extraction) + if model_ema: + avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( + model_ema, criterion, data_loader_test, device=device, transform=None, + log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op, + nn_for_feature_extraction=args.nn_for_feature_extraction) + if args.output_dir and avg_accuracy >= best['accuracy']: + logger.info(f"Epoch {epoch}: {avg_accuracy:.2f} (Val accuracy) >= {best['accuracy']:.2f} (So far best accuracy). Hence updating checkpoint.pth") + best['accuracy'], best['f1'], best['auc'], best['conf_matrix'], best['epoch'] = avg_accuracy, avg_f1, auc, avg_conf_matrix, epoch + best['predictions'], best['ground_truth'] = predictions, ground_truth + checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) + utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) + + if not args.quantization and args.auto_quantization: + _float_best_metric = best['accuracy'] / 100.0 + logger.info(f"Stored float best accuracy for binary search: {_float_best_metric:.4f}") + + # Log best epoch results + logger = getLogger(f"root.main.{phase}.BestEpoch") + logger.info("") + logger.info("Printing statistics of best epoch:") + logger.info(f"Best Epoch: {best['epoch']}") + logger.info(f"Acc@1 {best['accuracy']:.3f}") + logger.info(f"F1-Score {best['f1']:.3f}") + logger.info(f"AUC ROC Score {best['f1']:.3f}") + logger.info("") + logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame(best['conf_matrix'], + columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], + index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), + headers="keys", tablefmt='grid'))) + + Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, + name="root.utils.print_file_level_classification_summary", + append_log=True if args.quantization else False, console_log=False) + getLogger("root.utils.print_file_level_classification_summary").propagate = False + utils.print_file_level_classification_summary(dataset_test, best['predictions'], best['ground_truth'], phase) + logger.info(f"Generated file-level classification summary in: {args.file_level_classification_log}") + + # Export model + logger.info('Exporting model after training.') + if args.distributed is False or (args.distributed is True and int(os.environ['LOCAL_RANK']) == 0): + if args.nn_for_feature_extraction: + example_input = next(iter(data_loader_test))[0] + input_shape = (1,) + dataset.X_raw.shape[1:] + else: + example_input = next(iter(data_loader_test))[1] + input_shape = (1,) + dataset.X.shape[1:] + utils.export_model( + model, input_shape=input_shape, output_dir=args.output_dir, opset_version=args.opset_version, + quantization=args.quantization, example_input=example_input, generic_model=args.generic_model, + remove_hooks_for_jit=True if (args.quantization_method == TinyMLQuantizationMethod.PTQ and args.quantization) else False) + + log_training_time(start_time) + + if args.gen_golden_vectors: + generate_golden_vector_dir(args.output_dir) + output_int = get_output_int_flag(args) + generate_golden_vectors(args.output_dir, dataset, output_int, args.generic_model, args.nn_for_feature_extraction) + + finally: + shutdown_data_loaders(data_loader, data_loader_test) def run(args): diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/test_onnx.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/test_onnx.py index d6f0029d..11b75057 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/test_onnx.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/test_onnx.py @@ -52,6 +52,7 @@ load_onnx_model, run_distributed_test, ) +from ..common.train_base import shutdown_data_loaders dataset_loader_dict = {'GenericTSDatasetForecasting': GenericTSDatasetForecasting} @@ -100,89 +101,93 @@ def main(gpu, args): data_loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=args.batch_size, sampler=test_sampler, num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn) + try: - logger.info(f"Loading ONNX model: {args.model_path}") - ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) - - predicted = torch.tensor([]).to(device, non_blocking=True) - ground_truth = torch.tensor([]).to(device, non_blocking=True) - - for _, batched_data, batched_target in data_loader_test: - batched_data = batched_data.to(device, non_blocking=True).float() - batched_target = batched_target.to(device, non_blocking=True).float() - if transform: - batched_data = transform(batched_data) - for data in batched_data: - predicted = torch.cat((predicted, torch.tensor( - ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0] - ).to(device))) - ground_truth = torch.cat((ground_truth, batched_target)) - - predicted = predicted.view_as(ground_truth) - - logger = getLogger("root.main.test_data") - for idx, item in enumerate(dataset_test.header_row): - for target_variable_name in item: - logger.info(f"Variable {target_variable_name}:") - logger.info(f" SMAPE of {target_variable_name} across all predicted timesteps: {utils.smape(ground_truth[:, :, idx], predicted[:, :, idx]):.2f}%") - logger.info(f" R² of {target_variable_name} across all predicted timesteps: {utils.get_r2_score(predicted[:, :, idx], ground_truth[:, :, idx]):.4f}") - - # Log timestep specific metrics - for step in range(args.forecast_horizon): - logger.info(f" Timestep {step + 1}:") - logger.info(f" SMAPE: {utils.smape(ground_truth[:, step, idx], predicted[:, step, idx]):.2f}%") - logger.info(f" R²: {utils.get_r2_score(predicted[:, step, idx], ground_truth[:, step, idx]):.4f}") - - # Save final predictions and create visualizations - if args.output_dir and ground_truth is not None: - results_dir = os.path.join(args.output_dir, 'test_results') - os.makedirs(results_dir, exist_ok=True) - - # Save predictions in CSV format - utils.save_forecasting_predictions_csv( - ground_truth, - predicted, - results_dir, - dataset_test.header_row, - args.forecast_horizon, - ) - - plots_dir = os.path.join(results_dir, 'prediction_plots') - os.makedirs(plots_dir, exist_ok=True) - - # Create scatter plots for each variable - for idx, item in enumerate(dataset_test.header_row): - for target_variable_name in item: - fig, axes = plt.subplots(int(np.ceil(args.forecast_horizon / 2)), 2, figsize=(12, 5)) - axes = axes.flatten() - for step in range(args.forecast_horizon): - step_targets = ground_truth[:, step, idx] - step_outputs = predicted[:, step, idx] - - step_smape = utils.smape(ground_truth[:, step, idx], predicted[:, step, idx]) - step_r2 = utils.get_r2_score(predicted[:, step, idx], ground_truth[:, step, idx]) + logger.info(f"Loading ONNX model: {args.model_path}") + ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) - # Convert to numpy for matplotlib plotting - step_targets_np = step_targets.detach().cpu().numpy() if isinstance(step_targets, torch.Tensor) else step_targets - step_outputs_np = step_outputs.detach().cpu().numpy() if isinstance(step_outputs, torch.Tensor) else step_outputs + predicted = torch.tensor([]).to(device, non_blocking=True) + ground_truth = torch.tensor([]).to(device, non_blocking=True) - # Scatter plot - ax = axes[step] - ax.scatter(step_targets_np, step_outputs_np, alpha=0.5, label='Predictions') + for _, batched_data, batched_target in data_loader_test: + batched_data = batched_data.to(device, non_blocking=True).float() + batched_target = batched_target.to(device, non_blocking=True).float() + if transform: + batched_data = transform(batched_data) + for data in batched_data: + predicted = torch.cat((predicted, torch.tensor( + ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0] + ).to(device))) + ground_truth = torch.cat((ground_truth, batched_target)) - # Add perfect prediction line - min_val = min(step_targets_np.min(), step_outputs_np.min()) - max_val = max(step_targets_np.max(), step_outputs_np.max()) - ax.plot([min_val, max_val], [min_val, max_val], 'k--', label='Perfect Prediction') + predicted = predicted.view_as(ground_truth) - ax.set_xlabel(f"Actual Variable {target_variable_name}") - ax.set_ylabel(f"Predicted Variable {target_variable_name}") - ax.set_title(f"{step + 1}-step ahead\nR² = {step_r2:.4f}, SMAPE = {step_smape:.2f}%") - ax.legend() + logger = getLogger("root.main.test_data") + for idx, item in enumerate(dataset_test.header_row): + for target_variable_name in item: + logger.info(f"Variable {target_variable_name}:") + logger.info(f" SMAPE of {target_variable_name} across all predicted timesteps: {utils.smape(ground_truth[:, :, idx], predicted[:, :, idx]):.2f}%") + logger.info(f" R² of {target_variable_name} across all predicted timesteps: {utils.get_r2_score(predicted[:, :, idx], ground_truth[:, :, idx]):.4f}") - plt.tight_layout() - plt.savefig(os.path.join(plots_dir, f'{target_variable_name}_predictions.png')) - plt.close() + # Log timestep specific metrics + for step in range(args.forecast_horizon): + logger.info(f" Timestep {step + 1}:") + logger.info(f" SMAPE: {utils.smape(ground_truth[:, step, idx], predicted[:, step, idx]):.2f}%") + logger.info(f" R²: {utils.get_r2_score(predicted[:, step, idx], ground_truth[:, step, idx]):.4f}") + + # Save final predictions and create visualizations + if args.output_dir and ground_truth is not None: + results_dir = os.path.join(args.output_dir, 'test_results') + os.makedirs(results_dir, exist_ok=True) + + # Save predictions in CSV format + utils.save_forecasting_predictions_csv( + ground_truth, + predicted, + results_dir, + dataset_test.header_row, + args.forecast_horizon, + ) + + plots_dir = os.path.join(results_dir, 'prediction_plots') + os.makedirs(plots_dir, exist_ok=True) + + # Create scatter plots for each variable + for idx, item in enumerate(dataset_test.header_row): + for target_variable_name in item: + fig, axes = plt.subplots(int(np.ceil(args.forecast_horizon / 2)), 2, figsize=(12, 5)) + axes = axes.flatten() + for step in range(args.forecast_horizon): + step_targets = ground_truth[:, step, idx] + step_outputs = predicted[:, step, idx] + + step_smape = utils.smape(ground_truth[:, step, idx], predicted[:, step, idx]) + step_r2 = utils.get_r2_score(predicted[:, step, idx], ground_truth[:, step, idx]) + + # Convert to numpy for matplotlib plotting + step_targets_np = step_targets.detach().cpu().numpy() if isinstance(step_targets, torch.Tensor) else step_targets + step_outputs_np = step_outputs.detach().cpu().numpy() if isinstance(step_outputs, torch.Tensor) else step_outputs + + # Scatter plot + ax = axes[step] + ax.scatter(step_targets_np, step_outputs_np, alpha=0.5, label='Predictions') + + # Add perfect prediction line + min_val = min(step_targets_np.min(), step_outputs_np.min()) + max_val = max(step_targets_np.max(), step_outputs_np.max()) + ax.plot([min_val, max_val], [min_val, max_val], 'k--', label='Perfect Prediction') + + ax.set_xlabel(f"Actual Variable {target_variable_name}") + ax.set_ylabel(f"Predicted Variable {target_variable_name}") + ax.set_title(f"{step + 1}-step ahead\nR² = {step_r2:.4f}, SMAPE = {step_smape:.2f}%") + ax.legend() + + plt.tight_layout() + plt.savefig(os.path.join(plots_dir, f'{target_variable_name}_predictions.png')) + plt.close() + + finally: + shutdown_data_loaders(data_loader_test) def run(args): diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py index 9539de6c..6930fa26 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_forecasting/train.py @@ -79,6 +79,7 @@ get_output_int_flag, load_onnx_for_inference, create_data_loaders, + shutdown_data_loaders, ) dataset_loader_dict = {'GenericTSDatasetForecasting': GenericTSDatasetForecasting} @@ -168,176 +169,180 @@ def main(gpu, args): logger.info("Loading data:") data_loader, data_loader_test = create_data_loaders(dataset, dataset_test, train_sampler, test_sampler, args, gpu) + try: + + logger.info("Creating model") + model = models.get_model( + args.model, variables, total_forecast_outputs, input_features=input_features, model_config=args.model_config, + model_spec=args.model_spec, + dual_op=args.dual_op) + + log_model_summary(model, args, variables, input_features, logger) + model = load_pretrained_weights(model, args, logger) + + # if output_int not set by user, then set it to default of task_type + if args.output_int == None: + args.output_int = False + + if handle_export_only(model, args, variables, input_features, logger): + return + + move_model_to_device(model, device, logger) + criterion = nn.HuberLoss() + + global _float_best_metric + sample_inputs = None + sample_targets = None + bsearch_float_metric = None + bsearch_example_inputs = None + if args.auto_quantization and args.quantization: + try: + sample_data_iter = iter(data_loader) + _, sample_data_fe, sample_targets_raw = next(sample_data_iter) + sample_inputs = sample_data_fe.float().to(device) + sample_targets = sample_targets_raw.float().to(device).reshape(sample_targets_raw.shape[0], -1) + logger.info("Obtained sample data for auto quantization analysis") + except Exception as e: + logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") + bsearch_float_metric = _float_best_metric + try: + bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) + except Exception as e: + logger.warning(f"Could not get example inputs for binary search: {e}") + + model = utils.quantization_wrapped_model( + model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, + args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, + calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, + eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, + task_type='forecasting', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, + autoquant_tolerance_forecasting=args.autoquant_tolerance_forecasting) + + optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) + resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) + + phase = 'QuantTrain' if args.quantization else 'FloatTrain' + logger.info("Start training") + start_time = timeit.default_timer() + + best_epoch_values = { + 'epoch': -1, + 'true_values': None, + 'predictions': None, + 'overall_smape': float('inf'), + } + + for epoch in range(args.start_epoch, args.epochs): + if args.distributed: + train_sampler.set_epoch(epoch) + + utils.train_one_epoch_forecasting( + model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, + print_freq=args.print_freq, phase=phase, num_classes=total_forecast_outputs, dual_op=args.dual_op, + is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False) + + if not (args.quantization_method in ['PTQ'] and args.quantization): + lr_scheduler.step() - logger.info("Creating model") - model = models.get_model( - args.model, variables, total_forecast_outputs, input_features=input_features, model_config=args.model_config, - model_spec=args.model_spec, - dual_op=args.dual_op) - - log_model_summary(model, args, variables, input_features, logger) - model = load_pretrained_weights(model, args, logger) - - # if output_int not set by user, then set it to default of task_type - if args.output_int == None: - args.output_int = False - - if handle_export_only(model, args, variables, input_features, logger): - return - - move_model_to_device(model, device, logger) - criterion = nn.HuberLoss() - - global _float_best_metric - sample_inputs = None - sample_targets = None - bsearch_float_metric = None - bsearch_example_inputs = None - if args.auto_quantization and args.quantization: - try: - sample_data_iter = iter(data_loader) - _, sample_data_fe, sample_targets_raw = next(sample_data_iter) - sample_inputs = sample_data_fe.float().to(device) - sample_targets = sample_targets_raw.float().to(device).reshape(sample_targets_raw.shape[0], -1) - logger.info("Obtained sample data for auto quantization analysis") - except Exception as e: - logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") - bsearch_float_metric = _float_best_metric - try: - bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) - except Exception as e: - logger.warning(f"Could not get example inputs for binary search: {e}") - - model = utils.quantization_wrapped_model( - model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, - args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, - calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, - eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, - task_type='forecasting', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, - autoquant_tolerance_forecasting=args.autoquant_tolerance_forecasting) - - optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) - model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) - resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) - - phase = 'QuantTrain' if args.quantization else 'FloatTrain' - logger.info("Start training") - start_time = timeit.default_timer() - - best_epoch_values = { - 'epoch': -1, - 'true_values': None, - 'predictions': None, - 'overall_smape': float('inf'), - } - - for epoch in range(args.start_epoch, args.epochs): - if args.distributed: - train_sampler.set_epoch(epoch) - - utils.train_one_epoch_forecasting( - model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, - print_freq=args.print_freq, phase=phase, num_classes=total_forecast_outputs, dual_op=args.dual_op, - is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False) - - if not (args.quantization_method in ['PTQ'] and args.quantization): - lr_scheduler.step() - - target_tensor, prediction_tensor, overall_smape = utils.evaluate_forecasting( - model, criterion, data_loader_test, device=device, transform=None, phase=phase, - num_classes=total_forecast_outputs, dual_op=args.dual_op) - - if model_ema: target_tensor, prediction_tensor, overall_smape = utils.evaluate_forecasting( - model_ema, criterion, data_loader_test, device=device, transform=None, - log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op) - - if overall_smape < best_epoch_values['overall_smape']: - best_epoch_values['overall_smape'] = overall_smape - best_epoch_values['epoch'] = epoch - best_epoch_values['true_values'] = target_tensor.clone() - best_epoch_values['predictions'] = prediction_tensor.clone() - - if args.output_dir: - checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema, - extra_data={'metrics': {'overall_smape': overall_smape}}) - utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) - - logger.info(f"Epoch {epoch}: Best Overall SMAPE across all variables across all predicted timesteps so far: {best_epoch_values['overall_smape']:.2f}% (Epoch {best_epoch_values['epoch']})") - - if not args.quantization and args.auto_quantization: - _float_best_metric = float(best_epoch_values['overall_smape']) - logger.info(f"Stored float best SMAPE for binary search: {_float_best_metric:.4f}") - - # Log best epoch metrics - logger = getLogger(f"root.main.{phase}.BestEpoch") - logger.info("Printing statistics of best epoch:") - logger.info(f"Best epoch:{best_epoch_values['epoch'] + 1}") - logger.info(f"Overall SMAPE across all variables: {best_epoch_values['overall_smape']:.2f}%") - logger.info("Per-Variable Metrics:") - - for idx, item in enumerate(dataset.header_row): - for target_variable_name in item: - logger.info(f" Variable {target_variable_name}:") - logger.info(f" SMAPE of {target_variable_name} across all predicted timesteps: {utils.smape(best_epoch_values['true_values'][:, :, idx], best_epoch_values['predictions'][:, :, idx]):.2f}%") - logger.info(f" R² of {target_variable_name} across all predicted timesteps: {utils.get_r2_score(best_epoch_values['predictions'][:, :, idx], best_epoch_values['true_values'][:, :, idx]):.4f}") - - for step in range(args.forecast_horizon): - logger.info(f" Timestep {step + 1}:") - logger.info(f" SMAPE: {utils.smape(best_epoch_values['true_values'][:, step, idx], best_epoch_values['predictions'][:, step, idx]):.2f}%") - logger.info(f" R²: {utils.get_r2_score(best_epoch_values['predictions'][:, step, idx], best_epoch_values['true_values'][:, step, idx]):.4f}") - - # Save final predictions and create visualizations for best epoch - if args.output_dir and best_epoch_values['true_values'] is not None: - results_dir = os.path.join(args.output_dir, f'best_epoch_{best_epoch_values["epoch"]}_results') - os.makedirs(results_dir, exist_ok=True) - - utils.save_forecasting_predictions_csv( - best_epoch_values['true_values'], - best_epoch_values['predictions'], - results_dir, - dataset.header_row, - args.forecast_horizon, - ) - - plots_dir = os.path.join(results_dir, 'prediction_plots') - os.makedirs(plots_dir, exist_ok=True) + model, criterion, data_loader_test, device=device, transform=None, phase=phase, + num_classes=total_forecast_outputs, dual_op=args.dual_op) + + if model_ema: + target_tensor, prediction_tensor, overall_smape = utils.evaluate_forecasting( + model_ema, criterion, data_loader_test, device=device, transform=None, + log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op) + + if overall_smape < best_epoch_values['overall_smape']: + best_epoch_values['overall_smape'] = overall_smape + best_epoch_values['epoch'] = epoch + best_epoch_values['true_values'] = target_tensor.clone() + best_epoch_values['predictions'] = prediction_tensor.clone() + + if args.output_dir: + checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema, + extra_data={'metrics': {'overall_smape': overall_smape}}) + utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) + + logger.info(f"Epoch {epoch}: Best Overall SMAPE across all variables across all predicted timesteps so far: {best_epoch_values['overall_smape']:.2f}% (Epoch {best_epoch_values['epoch']})") + + if not args.quantization and args.auto_quantization: + _float_best_metric = float(best_epoch_values['overall_smape']) + logger.info(f"Stored float best SMAPE for binary search: {_float_best_metric:.4f}") + + # Log best epoch metrics + logger = getLogger(f"root.main.{phase}.BestEpoch") + logger.info("Printing statistics of best epoch:") + logger.info(f"Best epoch:{best_epoch_values['epoch'] + 1}") + logger.info(f"Overall SMAPE across all variables: {best_epoch_values['overall_smape']:.2f}%") + logger.info("Per-Variable Metrics:") for idx, item in enumerate(dataset.header_row): for target_variable_name in item: - fig, axes = plt.subplots(int(np.ceil(args.forecast_horizon / 2)), 2, figsize=(12, 5)) - axes = axes.flatten() + logger.info(f" Variable {target_variable_name}:") + logger.info(f" SMAPE of {target_variable_name} across all predicted timesteps: {utils.smape(best_epoch_values['true_values'][:, :, idx], best_epoch_values['predictions'][:, :, idx]):.2f}%") + logger.info(f" R² of {target_variable_name} across all predicted timesteps: {utils.get_r2_score(best_epoch_values['predictions'][:, :, idx], best_epoch_values['true_values'][:, :, idx]):.4f}") + for step in range(args.forecast_horizon): - step_targets = best_epoch_values['true_values'][:, step, idx] - step_outputs = best_epoch_values['predictions'][:, step, idx] - step_smape = utils.smape(best_epoch_values['true_values'][:, step, idx], best_epoch_values['predictions'][:, step, idx]) - step_r2 = utils.get_r2_score(best_epoch_values['predictions'][:, step, idx], best_epoch_values['true_values'][:, step, idx]) - - # Convert to numpy for matplotlib plotting - step_targets_np = step_targets.detach().cpu().numpy() if isinstance(step_targets, torch.Tensor) else step_targets - step_outputs_np = step_outputs.detach().cpu().numpy() if isinstance(step_outputs, torch.Tensor) else step_outputs - - ax = axes[step] - ax.scatter(step_targets_np, step_outputs_np, alpha=0.5, label='Predictions') - min_val = min(step_targets_np.min(), step_outputs_np.min()) - max_val = max(step_targets_np.max(), step_outputs_np.max()) - ax.plot([min_val, max_val], [min_val, max_val], 'k--', label='Perfect Prediction') - ax.set_xlabel(f"Actual Variable {target_variable_name}") - ax.set_ylabel(f"Predicted Variable {target_variable_name}") - ax.set_title(f"{step + 1}-step ahead\nR² = {step_r2:.4f},SMAPE = {step_smape:.2f}%") - ax.legend() - - plt.tight_layout() - plt.savefig(os.path.join(plots_dir, f'{target_variable_name}_predictions.png')) - plt.close() - - export_trained_model(model, args, dataset) - log_training_time(start_time) - - if args.gen_golden_vectors: - generate_golden_vector_dir(args.output_dir) - output_int = get_output_int_flag(args) - generate_golden_vectors(args.output_dir, output_int, dataset, args.generic_model) + logger.info(f" Timestep {step + 1}:") + logger.info(f" SMAPE: {utils.smape(best_epoch_values['true_values'][:, step, idx], best_epoch_values['predictions'][:, step, idx]):.2f}%") + logger.info(f" R²: {utils.get_r2_score(best_epoch_values['predictions'][:, step, idx], best_epoch_values['true_values'][:, step, idx]):.4f}") + + # Save final predictions and create visualizations for best epoch + if args.output_dir and best_epoch_values['true_values'] is not None: + results_dir = os.path.join(args.output_dir, f'best_epoch_{best_epoch_values["epoch"]}_results') + os.makedirs(results_dir, exist_ok=True) + + utils.save_forecasting_predictions_csv( + best_epoch_values['true_values'], + best_epoch_values['predictions'], + results_dir, + dataset.header_row, + args.forecast_horizon, + ) + + plots_dir = os.path.join(results_dir, 'prediction_plots') + os.makedirs(plots_dir, exist_ok=True) + + for idx, item in enumerate(dataset.header_row): + for target_variable_name in item: + fig, axes = plt.subplots(int(np.ceil(args.forecast_horizon / 2)), 2, figsize=(12, 5)) + axes = axes.flatten() + for step in range(args.forecast_horizon): + step_targets = best_epoch_values['true_values'][:, step, idx] + step_outputs = best_epoch_values['predictions'][:, step, idx] + step_smape = utils.smape(best_epoch_values['true_values'][:, step, idx], best_epoch_values['predictions'][:, step, idx]) + step_r2 = utils.get_r2_score(best_epoch_values['predictions'][:, step, idx], best_epoch_values['true_values'][:, step, idx]) + + # Convert to numpy for matplotlib plotting + step_targets_np = step_targets.detach().cpu().numpy() if isinstance(step_targets, torch.Tensor) else step_targets + step_outputs_np = step_outputs.detach().cpu().numpy() if isinstance(step_outputs, torch.Tensor) else step_outputs + + ax = axes[step] + ax.scatter(step_targets_np, step_outputs_np, alpha=0.5, label='Predictions') + min_val = min(step_targets_np.min(), step_outputs_np.min()) + max_val = max(step_targets_np.max(), step_outputs_np.max()) + ax.plot([min_val, max_val], [min_val, max_val], 'k--', label='Perfect Prediction') + ax.set_xlabel(f"Actual Variable {target_variable_name}") + ax.set_ylabel(f"Predicted Variable {target_variable_name}") + ax.set_title(f"{step + 1}-step ahead\nR² = {step_r2:.4f},SMAPE = {step_smape:.2f}%") + ax.legend() + + plt.tight_layout() + plt.savefig(os.path.join(plots_dir, f'{target_variable_name}_predictions.png')) + plt.close() + + export_trained_model(model, args, dataset) + log_training_time(start_time) + + if args.gen_golden_vectors: + generate_golden_vector_dir(args.output_dir) + output_int = get_output_int_flag(args) + generate_golden_vectors(args.output_dir, output_int, dataset, args.generic_model) + + finally: + shutdown_data_loaders(data_loader, data_loader_test) def run(args): diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/test_onnx.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/test_onnx.py index c7c77606..7fcfc260 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/test_onnx.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/test_onnx.py @@ -50,6 +50,7 @@ load_onnx_model, run_distributed_test, ) +from ..common.train_base import shutdown_data_loaders dataset_loader_dict = {'GenericTSDataset': GenericTSDataset, 'GenericTSDatasetReg': GenericTSDatasetReg} @@ -97,48 +98,51 @@ def main(gpu, args): data_loader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn) - - logger.info(f"Loading ONNX model: {args.model_path}") - ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) - - predicted = torch.tensor([]).to(device, non_blocking=True) - ground_truth = torch.tensor([]).to(device, non_blocking=True) - - for _, batched_data, batched_target in data_loader: - batched_data = batched_data.to(device, non_blocking=True).float() - batched_target = batched_target.to(device, non_blocking=True).float() - if transform: - batched_data = transform(batched_data) - for data in batched_data: - predicted = torch.cat((predicted, torch.tensor( - ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0] - ).to(device))) - ground_truth = torch.cat((ground_truth, batched_target)) - - mdcl_utils.create_dir(os.path.join(args.output_dir, 'post_training_analysis')) - logger.info("Plotting Regressions on dataset") - - metric = torcheval.metrics.MeanSquaredError() - r2_score = torcheval.metrics.R2Score() - - df = pd.DataFrame({ - "predicted": predicted.to('cpu').numpy().flatten(), - "ground_truth": ground_truth.to('cpu').numpy().flatten() - }) - df.to_csv(os.path.join(args.output_dir, 'post_training_analysis', "results_on_test_set.csv"), index=False) - logger.info(f"Outputs on the test set saved at : {os.path.join(args.output_dir, 'post_training_analysis', 'results_on_test_set.csv')}") - - utils.plot_actual_vs_predicted_regression(ground_truth.to('cpu'), predicted.to('cpu'), - os.path.join(args.output_dir, 'post_training_analysis'), phase='test') - utils.plot_residual_error_regression(ground_truth.to('cpu'), predicted.to('cpu'), - os.path.join(args.output_dir, 'post_training_analysis'), phase='test') - - metric.update(predicted.to('cpu'), ground_truth.to('cpu')) - r2_score.update(predicted.to('cpu'), ground_truth.to('cpu')) - - logger = getLogger("root.main.test_data") - logger.info(f"{logger.name}: Test Data Evaluation RMSE: {torch.sqrt(metric.compute()):.2f}") - logger.info(f"{logger.name}: Test Data Evaluation R2-Score: {r2_score.compute():.2f}") + try: + + logger.info(f"Loading ONNX model: {args.model_path}") + ort_sess, input_name, output_name = load_onnx_model(args.model_path, args.generic_model) + + predicted = torch.tensor([]).to(device, non_blocking=True) + ground_truth = torch.tensor([]).to(device, non_blocking=True) + + for _, batched_data, batched_target in data_loader: + batched_data = batched_data.to(device, non_blocking=True).float() + batched_target = batched_target.to(device, non_blocking=True).float() + if transform: + batched_data = transform(batched_data) + for data in batched_data: + predicted = torch.cat((predicted, torch.tensor( + ort_sess.run([output_name], {input_name: data.unsqueeze(0).cpu().numpy()})[0] + ).to(device))) + ground_truth = torch.cat((ground_truth, batched_target)) + + mdcl_utils.create_dir(os.path.join(args.output_dir, 'post_training_analysis')) + logger.info("Plotting Regressions on dataset") + + metric = torcheval.metrics.MeanSquaredError() + r2_score = torcheval.metrics.R2Score() + + df = pd.DataFrame({ + "predicted": predicted.to('cpu').numpy().flatten(), + "ground_truth": ground_truth.to('cpu').numpy().flatten() + }) + df.to_csv(os.path.join(args.output_dir, 'post_training_analysis', "results_on_test_set.csv"), index=False) + logger.info(f"Outputs on the test set saved at : {os.path.join(args.output_dir, 'post_training_analysis', 'results_on_test_set.csv')}") + + utils.plot_actual_vs_predicted_regression(ground_truth.to('cpu'), predicted.to('cpu'), + os.path.join(args.output_dir, 'post_training_analysis'), phase='test') + utils.plot_residual_error_regression(ground_truth.to('cpu'), predicted.to('cpu'), + os.path.join(args.output_dir, 'post_training_analysis'), phase='test') + + metric.update(predicted.to('cpu'), ground_truth.to('cpu')) + r2_score.update(predicted.to('cpu'), ground_truth.to('cpu')) + + logger = getLogger("root.main.test_data") + logger.info(f"{logger.name}: Test Data Evaluation RMSE: {torch.sqrt(metric.compute()):.2f}") + logger.info(f"{logger.name}: Test Data Evaluation R2-Score: {r2_score.compute():.2f}") + finally: + shutdown_data_loaders(data_loader) return diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/train.py index f003f945..7e23cf5c 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_regression/train.py @@ -76,6 +76,7 @@ get_output_int_flag, load_onnx_for_inference, create_data_loaders, + shutdown_data_loaders, ) dataset_loader_dict = {'GenericTSDatasetReg': GenericTSDatasetReg} @@ -171,98 +172,102 @@ def main(gpu, args): logger.info("Loading data:") data_loader, data_loader_test = create_data_loaders(dataset, dataset_test, train_sampler, test_sampler, args, gpu) + try: - logger.info("Creating model") - model = create_model(args, variables, num_classes, input_features, logger) - log_model_summary(model, args, variables, input_features, logger) - model = load_pretrained_weights(model, args, logger) - - # if output_int not set by user, then set it to default of task_type - if args.output_int == None: - args.output_int = False - - move_model_to_device(model, device, logger) - global _float_best_metric - sample_inputs = None - sample_targets = None - bsearch_float_metric = None - bsearch_example_inputs = None - if args.auto_quantization and args.quantization: - try: - sample_data_iter = iter(data_loader) - _, sample_data_fe, sample_targets_raw = next(sample_data_iter) - sample_inputs = sample_data_fe.float().to(device) - sample_targets = sample_targets_raw.float().to(device) - logger.info("Obtained sample data for auto quantization analysis") - except Exception as e: - logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") - bsearch_float_metric = _float_best_metric - try: - bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) - except Exception as e: - logger.warning(f"Could not get example inputs for binary search: {e}") - criterion = nn.MSELoss().to(device) - model = utils.quantization_wrapped_model( - model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, - args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, - calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, - eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, - task_type='regression', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, - autoquant_tolerance_regression=args.autoquant_tolerance_regression) - - if handle_export_only(model, args, variables, input_features, logger): - return - - optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) - model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) - resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) - - phase = 'QuantTrain' if args.quantization else 'FloatTrain' - logger.info("Start training") - start_time = timeit.default_timer() - best = dict(mse=np.inf, r2=0, epoch=None) - - for epoch in range(args.start_epoch, args.epochs): - if args.distributed: - train_sampler.set_epoch(epoch) - utils.train_one_epoch_regression( - model, criterion, optimizer, data_loader, device, epoch, None, args.lambda_reg, args.apex, model_ema, - print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, - is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False) - if not (args.quantization_method in ['PTQ'] and args.quantization): - lr_scheduler.step() - avg_mse, avg_r2_score = utils.evaluate_regression(model, criterion, data_loader_test, device=device, - transform=None, phase=phase, num_classes=num_classes, dual_op=args.dual_op) - if model_ema: - avg_mse, avg_r2_score = utils.evaluate_regression( - model_ema, criterion, data_loader_test, device=device, transform=None, - log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op) - if args.output_dir and avg_mse <= best['mse']: - logger.info(f"Epoch {epoch}: {avg_mse:.2f} (Val MSE) <= {best['mse']:.2f} (So far least error). Hence updating checkpoint.pth") - best['mse'], best['r2'], best['epoch'] = avg_mse, avg_r2_score, epoch - checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) - utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) - - if not args.quantization and args.auto_quantization: - _float_best_metric = best['r2'] - logger.info(f"Stored float best R² for binary search: {_float_best_metric:.4f}") - - # Log best epoch results - logger = getLogger(f"root.main.{phase}.BestEpoch") - logger.info("") - logger.info("Printing statistics of best epoch:") - logger.info(f"Best Epoch: {best['epoch']}") - logger.info(f"MSE {best['mse']:.3f}") - logger.info(f"R2-Score {best['r2']:.3f}") - logger.info("") - - export_trained_model(model, args, dataset) - log_training_time(start_time) - - if args.gen_golden_vectors: - generate_golden_vector_dir(args.output_dir) - output_int = get_output_int_flag(args) - generate_golden_vectors(args.output_dir, output_int, dataset, args.generic_model) + logger.info("Creating model") + model = create_model(args, variables, num_classes, input_features, logger) + log_model_summary(model, args, variables, input_features, logger) + model = load_pretrained_weights(model, args, logger) + + # if output_int not set by user, then set it to default of task_type + if args.output_int == None: + args.output_int = False + + move_model_to_device(model, device, logger) + global _float_best_metric + sample_inputs = None + sample_targets = None + bsearch_float_metric = None + bsearch_example_inputs = None + if args.auto_quantization and args.quantization: + try: + sample_data_iter = iter(data_loader) + _, sample_data_fe, sample_targets_raw = next(sample_data_iter) + sample_inputs = sample_data_fe.float().to(device) + sample_targets = sample_targets_raw.float().to(device) + logger.info("Obtained sample data for auto quantization analysis") + except Exception as e: + logger.warning(f"Could not obtain sample data for auto quantization: {e}. Proceeding without it.") + bsearch_float_metric = _float_best_metric + try: + bsearch_example_inputs = next(iter(data_loader_test))[1][:1].float().to(device) + except Exception as e: + logger.warning(f"Could not get example inputs for binary search: {e}") + criterion = nn.MSELoss().to(device) + model = utils.quantization_wrapped_model( + model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, + args.epochs, args.output_int, args.auto_quantization, inputs=sample_inputs, targets=sample_targets, criterion=criterion, + calibration_dataloader=data_loader if (args.auto_quantization and args.quantization) else None, + eval_dataloader=data_loader_test if (args.auto_quantization and args.quantization) else None, + task_type='regression', float_metric=bsearch_float_metric, example_inputs=bsearch_example_inputs, + autoquant_tolerance_regression=args.autoquant_tolerance_regression) + + if handle_export_only(model, args, variables, input_features, logger): + return + + optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) + resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) + + phase = 'QuantTrain' if args.quantization else 'FloatTrain' + logger.info("Start training") + start_time = timeit.default_timer() + best = dict(mse=np.inf, r2=0, epoch=None) + + for epoch in range(args.start_epoch, args.epochs): + if args.distributed: + train_sampler.set_epoch(epoch) + utils.train_one_epoch_regression( + model, criterion, optimizer, data_loader, device, epoch, None, args.lambda_reg, args.apex, model_ema, + print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, + is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False) + if not (args.quantization_method in ['PTQ'] and args.quantization): + lr_scheduler.step() + avg_mse, avg_r2_score = utils.evaluate_regression(model, criterion, data_loader_test, device=device, + transform=None, phase=phase, num_classes=num_classes, dual_op=args.dual_op) + if model_ema: + avg_mse, avg_r2_score = utils.evaluate_regression( + model_ema, criterion, data_loader_test, device=device, transform=None, + log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op) + if args.output_dir and avg_mse <= best['mse']: + logger.info(f"Epoch {epoch}: {avg_mse:.2f} (Val MSE) <= {best['mse']:.2f} (So far least error). Hence updating checkpoint.pth") + best['mse'], best['r2'], best['epoch'] = avg_mse, avg_r2_score, epoch + checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) + utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) + + if not args.quantization and args.auto_quantization: + _float_best_metric = best['r2'] + logger.info(f"Stored float best R² for binary search: {_float_best_metric:.4f}") + + # Log best epoch results + logger = getLogger(f"root.main.{phase}.BestEpoch") + logger.info("") + logger.info("Printing statistics of best epoch:") + logger.info(f"Best Epoch: {best['epoch']}") + logger.info(f"MSE {best['mse']:.3f}") + logger.info(f"R2-Score {best['r2']:.3f}") + logger.info("") + + export_trained_model(model, args, dataset) + log_training_time(start_time) + + if args.gen_golden_vectors: + generate_golden_vector_dir(args.output_dir) + output_int = get_output_int_flag(args) + generate_golden_vectors(args.output_dir, output_int, dataset, args.generic_model) + + finally: + shutdown_data_loaders(data_loader, data_loader_test) def run(args):