diff --git a/src/sparseml/pytorch/image_classification/lr_analysis.py b/src/sparseml/pytorch/image_classification/lr_analysis.py index 9782951c968..980d73ae6ab 100644 --- a/src/sparseml/pytorch/image_classification/lr_analysis.py +++ b/src/sparseml/pytorch/image_classification/lr_analysis.py @@ -13,83 +13,90 @@ # limitations under the License. """ -##### -Command help: -usage: lr_analysis.py [-h] --batch-size BATCH_SIZE --arch-key ARCH_KEY - --dataset DATASET --dataset-path DATASET_PATH [ - --pretrained PRETRAINED] - [--pretrained-dataset PRETRAINED_DATASET] [ - --model-kwargs MODEL_KWARGS] - [--dataset-kwargs DATASET_KWARGS] [--model-tag MODEL_TAG] - [--save-dir SAVE_DIR] [--device DEVICE] - [--loader-num-workers LOADER_NUM_WORKERS] [ - --no-loader-pin-memory] - [--loader-pin-memory [LOADER_PIN_MEMORY]] - [--checkpoint-path CHECKPOINT_PATH] [--init-lr INIT_LR] - [--optim-args OPTIM_ARGS] [--final-lr FINAL_LR] - [--steps-per-measurement STEPS_PER_MEASUREMENT] +Usage: sparseml.image_classification.lr_analysis [OPTIONS] + + Run a learning rate sensitivity analysis for a desired image classification + architecture + +Options: + --batch-size, --batch_size INTEGER + The batch size to use for analysis + [required] + -d, --dataset TEXT The dataset used for training, ex: + `imagenet`, `imagenette`, `cifar10`, etc. + Set to `imagefolder` for a generic dataset + setup with imagefolder type structure like + imagenet or loadable by a dataset in + `sparseml.pytorch.datasets` [required] + --dataset-path, --dataset_path DIRECTORY + The root dir path where the dataset is + stored or should be downloaded to if + available [required] + --arch_key, --arch-key TEXT The architecture key for image + classification model; example: `resnet50`, + `mobilenet`. Note: Will be read from the + checkpoint if not specified + --pretrained TEXT The type of pretrained weights to use, loads + default pretrained weights for the model if + not specified or set to `True`. Otherwise + should be set to the desired weights type: + [base, optim, optim-perf]. To not load any + weights set to one of [none, false] + [default: True] + --pretrained-dataset, --pretrained_dataset TEXT + The dataset to load pretrained weights for + if pretrained is set. Load the default + dataset for the architecture if set to None. + examples:`imagenet`, `cifar10`, etc... + --model-kwargs, --model_kwargs TEXT + Keyword arguments to be passed to model + constructor, should be given as a json + object + --dataset-kwargs, --dataset_kwargs TEXT + Keyword arguments to be passed to dataset + constructor, should be specified as a json + object + --model-tag, --model_tag TEXT A tag for saving results under save-dir, + defaults to the model arch and dataset used + --save-dir, --save_dir DIRECTORY + The path to the directory for saving results + [default: pytorch_vision] + --device TEXT The device to run on (can also include ids + for data parallel), ex: cpu, cuda, cuda:0,1 + [default: cuda] + --loader-num-workers, --loader_num_workers INTEGER + The number of workers to use for data + loading + --loader-pin-memory, --loader_pin_memory / --loader-no-pin-memory, + --loader_no_pin_memory + Use pinned memory for data loading + [default: loader-pin-memory] + --checkpoint-path, --checkpoint_path TEXT + A path to a previous checkpoint to load the + state from and resume the state for. If + provided, pretrained will be ignored . If + using a SparseZoo recipe, can also provide + 'zoo' to load the base weights associated + with that recipe. Additionally, can also + provide a SparseZoo model stub to load model + weights from SparseZoo + --init-lr, --init_lr FLOAT The initial learning rate to use for + analysis [default: 1e-09] + --optim-args, --optimizer-args, --optim_args, --optimizer_args TEXT + Additional args to be passed to the + optimizer; should be specified as a json + object + --final-lr, --final_lr FLOAT The final learning rate to use for the + sensitivity analysis [default: 0.5] + --steps-per-measurement, --steps_per_measurement INTEGER + The number of steps (batches) to run for + each measurement + -is, --image-size, --image_size INTEGER + The size of the image input to the model. + Value should be equal to S for [C, S, S] or + [S, S, C] dimensional input [default: 224] + --help Show this message and exit. -Utility script to Run a learning rate sensitivity analysis for a desired -image classification architecture - -optional arguments: - -h, --help show this help message and exit - --batch-size BATCH_SIZE - The batch size to use for analysis - --arch-key ARCH_KEY The type of model to use, ex: resnet50, vgg16, mobilenet - put as help to see the full list (will raise an - exception with the list) - --dataset DATASET The dataset to use for analysis, ex: imagenet, - imagenette, cifar10, - etc. Set to imagefolder for a generic dataset setup - with an imagefolder structure setup like imagenet or - loadable by a dataset in sparseml.pytorch.datasets - --dataset-path DATASET_PATH - The root path to where the dataset is stored - --pretrained PRETRAINED - The type of pretrained weights to use, default is - true to load the default pretrained weights for the - model. Otherwise should be set to the desired weights - type: [base, optim, optim-perf]. To not load any - weights set to one of [none, false] - --pretrained-dataset PRETRAINED_DATASET - The dataset to load pretrained weights for if pretrained - is set. Default is None which will load the default - dataset for the architecture. Ex can be set to - imagenet, cifar10, etc - --model-kwargs MODEL_KWARGS - Keyword arguments to be passed to model constructor, - should be given as a json object - --dataset-kwargs DATASET_KWARGS - Keyword arguments to be passed to dataset - constructor, should be given as a json object - --model-tag MODEL_TAG - A tag to use for the model for saving results under - save-dir, defaults to the model arch and dataset used - --save-dir SAVE_DIR The path to the directory for saving results - --device DEVICE The device to run on (can also include ids for - data parallel), ex: cpu, cuda, cuda:0,1 - --loader-num-workers LOADER_NUM_WORKERS - The number of workers to use for data loading - --no-loader-pin-memory - Do not use pinned memory for data loading - --loader-pin-memory [LOADER_PIN_MEMORY] - Use pinned memory for data loading - --checkpoint-path CHECKPOINT_PATH - A path to a previous checkpoint to load the state from - and resume the state for. If provided, pretrained - will be ignored. If using a SparseZoo recipe, - can also provide 'zoo' to load the base weights - associated with that recipe - --init-lr INIT_LR The initial learning rate to use for the sensitivity - analysis - --optim-args OPTIM_ARGS - Additional args to be passed to the optimizer passed in - as a json object - --final-lr FINAL_LR The final learning rate to use for the sensitivity - analysis - --steps-per-measurement STEPS_PER_MEASUREMENT The number of steps (batches) - to run for each measurement ######### EXAMPLES ######### @@ -105,20 +112,19 @@ sparseml.image_classification.lr_analysis \ --arch-key mobilenet --dataset imagenette \ --dataset-path ~/datasets/ --batch-size 2 -""" +""" import json import os -from dataclasses import dataclass, field -from typing import Optional +from typing import Any, Dict, Optional, Union from torch.nn import Module from torch.optim import SGD from torch.utils.data import DataLoader +import click from sparseml import get_main_logger -from sparseml.pytorch.image_classification.utils import NmArgumentParser, helpers -from sparseml.pytorch.models import ModelRegistry +from sparseml.pytorch.image_classification.utils import helpers from sparseml.pytorch.optim import default_exponential_check_lrs, lr_loss_sensitivity from sparseml.pytorch.utils import ( CrossEntropyLossWrapper, @@ -132,223 +138,279 @@ LOGGER = get_main_logger() -@dataclass -class LRAnalysisArguments: +@click.command() +@click.option( + "--batch-size", + "--batch_size", + type=int, + required=True, + help="The batch size to use for analysis", +) +@click.option( + "--dataset", + "-d", + type=str, + required=True, + help="The dataset used for training, " + "ex: `imagenet`, `imagenette`, `cifar10`, etc. " + "Set to `imagefolder` for a generic dataset setup with " + "imagefolder type structure like imagenet or loadable by " + "a dataset in `sparseml.pytorch.datasets`", +) +@click.option( + "--dataset-path", + "--dataset_path", + type=click.Path(dir_okay=True, file_okay=False), + callback=helpers.create_dir_callback, + required=True, + help="The root dir path where the dataset is stored or should " + "be downloaded to if available", +) +@click.option( + "--arch_key", + "--arch-key", + type=str, + default=None, + help="The architecture key for image classification model; " + "example: `resnet50`, `mobilenet`. " + "Note: Will be read from the checkpoint if not specified", +) +@click.option( + "--pretrained", + type=str, + default=True, + show_default=True, + help="The type of pretrained weights to use, " + "loads default pretrained weights for " + "the model if not specified or set to `True`. " + "Otherwise should be set to the desired weights " + "type: [base, optim, optim-perf]. To not load any weights set" + " to one of [none, false]", +) +@click.option( + "--pretrained-dataset", + "--pretrained_dataset", + type=str, + default=None, + show_default=True, + help="The dataset to load pretrained weights for if pretrained is " + "set. Load the default dataset for the architecture if set to None. " + "examples:`imagenet`, `cifar10`, etc...", +) +@click.option( + "--model-kwargs", + "--model_kwargs", + default=json.dumps({}), + type=str, + callback=helpers.parse_json_callback, + help="Keyword arguments to be passed to model constructor, should " + "be given as a json object", +) +@click.option( + "--dataset-kwargs", + "--dataset_kwargs", + default=json.dumps({}), + type=str, + callback=helpers.parse_json_callback, + help="Keyword arguments to be passed to dataset constructor, " + "should be specified as a json object", +) +@click.option( + "--model-tag", + "--model_tag", + type=str, + default=None, + help="A tag for saving results under save-dir, " + "defaults to the model arch and dataset used", +) +@click.option( + "--save-dir", + "--save_dir", + type=click.Path(dir_okay=True, file_okay=False), + default="pytorch_vision", + callback=helpers.create_dir_callback, + show_default=True, + help="The path to the directory for saving results", +) +@click.option( + "--device", + default=default_device(), + show_default=True, + help="The device to run on (can also include ids for data " + "parallel), ex: cpu, cuda, cuda:0,1", +) +@click.option( + "--loader-num-workers", + "--loader_num_workers", + type=int, + default=4, + help="The number of workers to use for data loading", +) +@click.option( + "--loader-pin-memory/--loader-no-pin-memory", + "--loader_pin_memory/--loader_no_pin_memory", + default=True, + is_flag=True, + show_default=True, + help="Use pinned memory for data loading", +) +@click.option( + "--checkpoint-path", + "--checkpoint_path", + type=str, + default=None, + help="A path to a previous checkpoint to load the state from " + "and resume the state for. If provided, pretrained will " + "be ignored . If using a SparseZoo recipe, can also " + "provide 'zoo' to load the base weights associated with " + "that recipe. Additionally, can also provide a SparseZoo model stub " + "to load model weights from SparseZoo", +) +@click.option( + "--init-lr", + "--init_lr", + type=float, + default=1e-9, + show_default=True, + help="The initial learning rate to use for analysis", +) +@click.option( + "--optim-args", + "--optimizer-args", + "--optim_args", + "--optimizer_args", + default=json.dumps({}), + type=str, + callback=helpers.parse_json_callback, + help="Additional args to be passed to the optimizer; " + "should be specified as a json object", +) +@click.option( + "--final-lr", + "--final_lr", + type=float, + default=0.5, + show_default=True, + help="The final learning rate to use for the sensitivity analysis", +) +@click.option( + "--steps-per-measurement", + "--steps_per_measurement", + type=int, + default=20, + help="The number of steps (batches) to run for each measurement", +) +@click.option( + "--image-size", + "--image_size", + "-is", + type=int, + default=224, + show_default=True, + help="The size of the image input to the model. Value should be " + "equal to S for [C, S, S] or [S, S, C] dimensional input", +) +def main( + batch_size: int, + dataset: str, + dataset_path: str, + pretrained: Union[str, bool], + pretrained_dataset: Optional[str], + model_kwargs: Dict[str, Any], + dataset_kwargs: Dict[str, Any], + arch_key: Optional[str], + checkpoint_path: Optional[str], + init_lr: float, + optim_args: Dict[str, Any], + final_lr: float, + model_tag: Optional[str], + save_dir: str, + device: Optional[str], + loader_num_workers: int, + loader_pin_memory: bool, + steps_per_measurement: int, + image_size: int, +): """ - Represents the arguments we use in our PyTorch integration scripts for - learning rate analysis. - Using :class:`NmArgumentParser` we can turn this class into `argparse - ` - arguments that can be specified on the command line. - - :param batch_size: The batch size to use for analysis. - :param arch_key: A str key representing the type of model to use, - ex:resnet50. - :param dataset: The dataset to use for analysis, ex imagenet, - imagenette, etc; Set to `imagefolder` for a custom dataset. - :param dataset_path: Root path to dataset location. - :param pretrained: The type of pretrained weights to use default is true to - load the default pretrained weights for the model. Otherwise should be - set to the desired weights type: [base, optim, optim-perf]; - To not load any weights set to one of [none, false]. - :param pretrained_dataset: str representing the dataset to load pretrained - weights for; if pretrained is set; Default is None which will load the - default dataset for the architecture; Ex can be set to imagenet, - cifar10, etc. - :param model_kwargs: json object containing keyword arguments to be passed - to the model constructor. - :param dataset_kwargs: json object to load keyword arguments to be passed - to dataset constructor. - :param model_tag: A str tag to use for the model for saving results under - save-dir, defaults to the model arch and dataset used. - :param save_dir: The path to the directory for saving results, - default="pytorch_vision". - :param device: str represnting the device to run on (can also include ids - for data parallel), ex:{cpu, cuda, cuda:0,1}. - :param loader_num_workers: int number of workers to use for data loading, - default=4. - :param loader_pin_memory: bool to use pinned memory for data loading, - default=True. - :param checkpoint_path: A path to a previous checkpoint to load the state - from and resume the state for; Also works with SparseZoo recipes; - Set to zoo to automatically download and load weights associated with a - recipe. - :param init_lr: float representing the initial learning for analysis, - default=1e-5. - :param optim_args: Additional arguments to be passed in to the optimizer as - a json object. - :param final_lr: The final learning rate to use for the sensitivity - analysis. - :param steps_per_measurement: The number of steps (batches) to run for each - measurement. + Run a learning rate sensitivity analysis for a desired + image classification architecture """ + is_main_process = True + local_rank = -1 - batch_size: int = field(metadata={"help": "The batch size to use for analysis"}) - - dataset: str = field( - metadata={ - "help": "The dataset to use for analysis, " - "ex: imagenet, imagenette, cifar10, etc. " - "Set to imagefolder for a generic dataset setup " - "with an image folder structure setup like imagenet or " - "loadable by a dataset in sparseml.pytorch.datasets" - } - ) - - dataset_path: str = field( - metadata={ - "help": "The root path to where the dataset is stored", - } - ) - pretrained: str = field( - default=True, - metadata={ - "help": "The type of pretrained weights to use, " - "default is true to load the default pretrained weights " - "for the model. " - "Otherwise should be set to the desired weights type: " - "[base, optim, optim-perf]. " - "To not load any weights set to one of [none, false]" - }, - ) - arch_key: Optional[str] = field( - default=None, - metadata={ - "help": "The type of model to use, ex: resnet50, vgg16, mobilenet " - "put as help to see the full list" - "(will raise an exception with the list)", - }, - ) - - pretrained_dataset: str = field( - default=None, - metadata={ - "help": "The dataset to load pretrained weights for if pretrained is" - "set. Default is None which will load the default dataset for " - "the architecture. Ex can be set to imagenet, cifar10, etc.", - }, - ) - - model_kwargs: json.loads = field( - default_factory=lambda: {}, - metadata={ - "help": "Keyword arguments to be passed to model constructor, should " - "be given as a json object" - }, - ) - - dataset_kwargs: json.loads = field( - default_factory=lambda: {}, - metadata={ - "help": "Keyword arguments to be passed to dataset constructor," - " should be given as a json object", - }, - ) - - model_tag: str = field( - default=None, - metadata={ - "help": "A tag to use for the model for saving results under save-dir, " - "defaults to the model arch and dataset used", - }, - ) - - save_dir: str = field( - default="pytorch_vision", - metadata={ - "help": "The path to the directory for saving results", - }, - ) - - device: str = field( - default=default_device(), - metadata={ - "help": "The device to run on (can also include ids for " - "data parallel), ex:cpu, cuda, cuda:0,1" - }, - ) - - loader_num_workers: int = field( - default=4, metadata={"help": "The number of workers to use for data loading"} - ) - - loader_pin_memory: bool = field( - default=True, metadata={"help": "Use pinned memory for data loading"} - ) - - checkpoint_path: str = field( - default=None, - metadata={ - "help": "A path to a previous checkpoint to load the state from " - "and resume the state for. If provided, pretrained will be" - "ignored. If using a SparseZoo recipe, can also provide " - "'zoo' to load the base weights associated with that recipe" - }, + save_dir, _ = helpers.get_save_dir_and_loggers( + task=CURRENT_TASK, + is_main_process=is_main_process, + save_dir=save_dir, + arch_key=arch_key, + model_tag=model_tag, + dataset_name=dataset, ) - init_lr: float = field( - default=1e-5, - metadata={ - "help": "The initial learning rate to use for the sensitivity analysis" - }, + train_dataset, train_loader = helpers.get_dataset_and_dataloader( + dataset_name=dataset, + dataset_path=dataset_path, + batch_size=batch_size, + image_size=image_size, + dataset_kwargs=dataset_kwargs, + training=True, + loader_num_workers=loader_num_workers, + loader_pin_memory=loader_pin_memory, ) - optim_args: json.loads = field( - default_factory=lambda: {}, - metadata={ - "help": "Additional args to be passed to the optimizer passed in" - " as a json object" - }, + num_classes = helpers.infer_num_classes( + train_dataset=train_dataset, + val_dataset=None, + dataset=dataset, + model_kwargs=model_kwargs, ) - - final_lr: float = field( - default=0.5, - metadata={ - "help": "The final learning rate to use for the sensitivity analysis", - }, + model, arch_key = helpers.create_model( + checkpoint_path=checkpoint_path, + recipe_path=None, + num_classes=num_classes, + arch_key=arch_key, + pretrained=pretrained, + pretrained_dataset=pretrained_dataset, + local_rank=local_rank, + **model_kwargs, ) - steps_per_measurement: int = field( - default=20, - metadata={"help": "The number of steps (batches) to run for each measurement"}, + lr_sensitivity( + model=model, + train_loader=train_loader, + save_dir=save_dir, + init_lr=init_lr, + optim_args=optim_args, + steps_per_measurement=steps_per_measurement, + device=device, + final_lr=final_lr, ) - def __post_init__(self): - self.arch_key = helpers.get_arch_key( - arch_key=self.arch_key, - checkpoint_path=self.checkpoint_path, - ) - - if "preprocessing_type" not in self.dataset_kwargs and ( - "coco" in self.dataset.lower() or "voc" in self.dataset.lower() - ): - if "ssd" in self.arch_key.lower(): - self.dataset_kwargs["preprocessing_type"] = "ssd" - elif "yolo" in self.arch_key.lower(): - self.dataset_kwargs["preprocessing_type"] = "yolo" - - self.is_main_process = True - self.local_rank = -1 - self.rank = -1 - def lr_sensitivity( - args, model: Module, train_loader: DataLoader, save_dir: str, + init_lr: float, + optim_args: Dict[str, Any], + steps_per_measurement: int, + device: Union[str, int], + final_lr: float, ) -> None: """ Utility function to run learning rate sensitivity analysis - :param args: An LRAnalysisArguments object containing config for current - LR analysis task. :param model: loaded model architecture to analyse :param train_loader: A DataLoader for training data :param save_dir: Directory to save results + :param init_lr: Initial learning rate to use for analysis + :param optim_args: Additional arguments to pass to the optimizer + :param steps_per_measurement: Number of steps to run for each measurement + :param device: Device to use for analysis + :param final_lr: Final learning rate to use for analysis """ # optimizer setup - optim = SGD(model.parameters(), lr=args.init_lr, **args.optim_args) + optim = SGD(model.parameters(), lr=init_lr, **optim_args) LOGGER.info(f"created optimizer: {optim}") # loss setup @@ -356,18 +418,18 @@ def lr_sensitivity( LOGGER.info(f"created loss: {loss}") # device setup - model, device, device_ids = model_to_device(model, args.device) + model, device, device_ids = model_to_device(model, device) # learning rate analysis LOGGER.info(f"running analysis: {loss}") analysis = lr_loss_sensitivity( - model, - train_loader, - loss, - optim, - device, - args.steps_per_measurement, - check_lrs=default_exponential_check_lrs(args.init_lr, args.final_lr), + module=model, + data=train_loader, + loss=loss, + optim=optim, + device=device, + steps_per_measurement=steps_per_measurement, + check_lrs=default_exponential_check_lrs(init_lr, final_lr), trainer_loggers=[PythonLogger()], ) @@ -379,59 +441,5 @@ def lr_sensitivity( analysis.print_res() -def main(): - """ - Driver function for the script - """ - _parser = NmArgumentParser( - dataclass_types=LRAnalysisArguments, - description="Utility script to Run a " - "learning rate sensitivity analysis " - "for a desired image classification architecture", - ) - args_, _ = _parser.parse_args_into_dataclasses() - save_dir, loggers = helpers.get_save_dir_and_loggers( - task=CURRENT_TASK, - is_main_process=args_.is_main_process, - save_dir=args_.save_dir, - arch_key=args_.arch_key, - model_tag=args_.model_tag, - dataset_name=args_.dataset, - ) - - input_shape = ModelRegistry.input_shape(args_.arch_key) - # assume shape [C, S, S] where S is the image size - image_size = input_shape[1] - - train_dataset, train_loader = helpers.get_dataset_and_dataloader( - dataset_name=args_.dataset, - dataset_path=args_.dataset_path, - batch_size=args_.batch_size, - image_size=image_size, - dataset_kwargs=args_.dataset_kwargs, - training=True, - loader_num_workers=args_.loader_num_workers, - loader_pin_memory=args_.loader_pin_memory, - ) - - num_classes = helpers.infer_num_classes( - train_dataset=train_dataset, - val_dataset=None, - dataset=args_.dataset, - model_kwargs=args_.model_kwargs, - ) - model, args_.arch_key = helpers.create_model( - checkpoint_path=args_.checkpoint_path, - recipe_path=None, - num_classes=num_classes, - arch_key=args_.arch_key, - pretrained=args_.pretrained, - pretrained_dataset=args_.pretrained_dataset, - local_rank=args_.local_rank, - **args_.model_kwargs, - ) - lr_sensitivity(args_, model, train_loader, save_dir) - - if __name__ == "__main__": main()