From dcad288fd6ff5109042c6fe61994db5cc5e55f3a Mon Sep 17 00:00:00 2001 From: Jakub Kaczmarzyk Date: Wed, 18 May 2022 10:27:33 -0400 Subject: [PATCH] use argparse groups to group arguments --- train.py | 227 +++++++++++++++++++++++++++++-------------------------- 1 file changed, 118 insertions(+), 109 deletions(-) diff --git a/train.py b/train.py index b6aade12..c953eb02 100755 --- a/train.py +++ b/train.py @@ -71,238 +71,247 @@ parser.add_argument('-c', '--config', default='', type=str, metavar='FILE', parser = argparse.ArgumentParser(description='PyTorch ImageNet Training') # Dataset parameters +group = parser.add_argument_group('Dataset parameters') +# Keep this argument outside of the dataset group because it is positional. parser.add_argument('data_dir', metavar='DIR', help='path to dataset') -parser.add_argument('--dataset', '-d', metavar='NAME', default='', +group.add_argument('--dataset', '-d', metavar='NAME', default='', help='dataset type (default: ImageFolder/ImageTar if empty)') -parser.add_argument('--train-split', metavar='NAME', default='train', +group.add_argument('--train-split', metavar='NAME', default='train', help='dataset train split (default: train)') -parser.add_argument('--val-split', metavar='NAME', default='validation', +group.add_argument('--val-split', metavar='NAME', default='validation', help='dataset validation split (default: validation)') -parser.add_argument('--dataset-download', action='store_true', default=False, +group.add_argument('--dataset-download', action='store_true', default=False, help='Allow download of dataset for torch/ and tfds/ datasets that support it.') -parser.add_argument('--class-map', default='', type=str, metavar='FILENAME', +group.add_argument('--class-map', default='', type=str, metavar='FILENAME', help='path to class to idx mapping file (default: "")') # Model parameters -parser.add_argument('--model', default='resnet50', type=str, metavar='MODEL', +group = parser.add_argument_group('Model parameters') +group.add_argument('--model', default='resnet50', type=str, metavar='MODEL', help='Name of model to train (default: "resnet50"') -parser.add_argument('--pretrained', action='store_true', default=False, +group.add_argument('--pretrained', action='store_true', default=False, help='Start with pretrained version of specified network (if avail)') -parser.add_argument('--initial-checkpoint', default='', type=str, metavar='PATH', +group.add_argument('--initial-checkpoint', default='', type=str, metavar='PATH', help='Initialize model from this checkpoint (default: none)') -parser.add_argument('--resume', default='', type=str, metavar='PATH', +group.add_argument('--resume', default='', type=str, metavar='PATH', help='Resume full model and optimizer state from checkpoint (default: none)') -parser.add_argument('--no-resume-opt', action='store_true', default=False, +group.add_argument('--no-resume-opt', action='store_true', default=False, help='prevent resume of optimizer state when resuming model') -parser.add_argument('--num-classes', type=int, default=None, metavar='N', +group.add_argument('--num-classes', type=int, default=None, metavar='N', help='number of label classes (Model default if None)') -parser.add_argument('--gp', default=None, type=str, metavar='POOL', +group.add_argument('--gp', default=None, type=str, metavar='POOL', help='Global pool type, one of (fast, avg, max, avgmax, avgmaxc). Model default if None.') -parser.add_argument('--img-size', type=int, default=None, metavar='N', +group.add_argument('--img-size', type=int, default=None, metavar='N', help='Image patch size (default: None => model default)') -parser.add_argument('--input-size', default=None, nargs=3, type=int, +group.add_argument('--input-size', default=None, nargs=3, type=int, metavar='N N N', help='Input all image dimensions (d h w, e.g. --input-size 3 224 224), uses model default if empty') -parser.add_argument('--crop-pct', default=None, type=float, +group.add_argument('--crop-pct', default=None, type=float, metavar='N', help='Input image center crop percent (for validation only)') -parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN', +group.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN', help='Override mean pixel value of dataset') -parser.add_argument('--std', type=float, nargs='+', default=None, metavar='STD', +group.add_argument('--std', type=float, nargs='+', default=None, metavar='STD', help='Override std deviation of dataset') -parser.add_argument('--interpolation', default='', type=str, metavar='NAME', +group.add_argument('--interpolation', default='', type=str, metavar='NAME', help='Image resize interpolation type (overrides model)') -parser.add_argument('-b', '--batch-size', type=int, default=128, metavar='N', +group.add_argument('-b', '--batch-size', type=int, default=128, metavar='N', help='Input batch size for training (default: 128)') -parser.add_argument('-vb', '--validation-batch-size', type=int, default=None, metavar='N', +group.add_argument('-vb', '--validation-batch-size', type=int, default=None, metavar='N', help='Validation batch size override (default: None)') -parser.add_argument('--channels-last', action='store_true', default=False, +group.add_argument('--channels-last', action='store_true', default=False, help='Use channels_last memory layout') -parser.add_argument('--torchscript', dest='torchscript', action='store_true', +group.add_argument('--torchscript', dest='torchscript', action='store_true', help='torch.jit.script the full model') -parser.add_argument('--fuser', default='', type=str, +group.add_argument('--fuser', default='', type=str, help="Select jit fuser. One of ('', 'te', 'old', 'nvfuser')") -parser.add_argument('--grad-checkpointing', action='store_true', default=False, +group.add_argument('--grad-checkpointing', action='store_true', default=False, help='Enable gradient checkpointing through model blocks/stages') # Optimizer parameters -parser.add_argument('--opt', default='sgd', type=str, metavar='OPTIMIZER', +group = parser.add_argument_group('Optimizer parameters') +group.add_argument('--opt', default='sgd', type=str, metavar='OPTIMIZER', help='Optimizer (default: "sgd"') -parser.add_argument('--opt-eps', default=None, type=float, metavar='EPSILON', +group.add_argument('--opt-eps', default=None, type=float, metavar='EPSILON', help='Optimizer Epsilon (default: None, use opt default)') -parser.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA', +group.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA', help='Optimizer Betas (default: None, use opt default)') -parser.add_argument('--momentum', type=float, default=0.9, metavar='M', +group.add_argument('--momentum', type=float, default=0.9, metavar='M', help='Optimizer momentum (default: 0.9)') -parser.add_argument('--weight-decay', type=float, default=2e-5, +group.add_argument('--weight-decay', type=float, default=2e-5, help='weight decay (default: 2e-5)') -parser.add_argument('--clip-grad', type=float, default=None, metavar='NORM', +group.add_argument('--clip-grad', type=float, default=None, metavar='NORM', help='Clip gradient norm (default: None, no clipping)') -parser.add_argument('--clip-mode', type=str, default='norm', +group.add_argument('--clip-mode', type=str, default='norm', help='Gradient clipping mode. One of ("norm", "value", "agc")') -parser.add_argument('--layer-decay', type=float, default=None, +group.add_argument('--layer-decay', type=float, default=None, help='layer-wise learning rate decay (default: None)') # Learning rate schedule parameters -parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER', +group = parser.add_argument_group('Learning rate schedule parameters') +group.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER', help='LR scheduler (default: "step"') -parser.add_argument('--lr', type=float, default=0.05, metavar='LR', +group.add_argument('--lr', type=float, default=0.05, metavar='LR', help='learning rate (default: 0.05)') -parser.add_argument('--lr-noise', type=float, nargs='+', default=None, metavar='pct, pct', +group.add_argument('--lr-noise', type=float, nargs='+', default=None, metavar='pct, pct', help='learning rate noise on/off epoch percentages') -parser.add_argument('--lr-noise-pct', type=float, default=0.67, metavar='PERCENT', +group.add_argument('--lr-noise-pct', type=float, default=0.67, metavar='PERCENT', help='learning rate noise limit percent (default: 0.67)') -parser.add_argument('--lr-noise-std', type=float, default=1.0, metavar='STDDEV', +group.add_argument('--lr-noise-std', type=float, default=1.0, metavar='STDDEV', help='learning rate noise std-dev (default: 1.0)') -parser.add_argument('--lr-cycle-mul', type=float, default=1.0, metavar='MULT', +group.add_argument('--lr-cycle-mul', type=float, default=1.0, metavar='MULT', help='learning rate cycle len multiplier (default: 1.0)') -parser.add_argument('--lr-cycle-decay', type=float, default=0.5, metavar='MULT', +group.add_argument('--lr-cycle-decay', type=float, default=0.5, metavar='MULT', help='amount to decay each learning rate cycle (default: 0.5)') -parser.add_argument('--lr-cycle-limit', type=int, default=1, metavar='N', +group.add_argument('--lr-cycle-limit', type=int, default=1, metavar='N', help='learning rate cycle limit, cycles enabled if > 1') -parser.add_argument('--lr-k-decay', type=float, default=1.0, +group.add_argument('--lr-k-decay', type=float, default=1.0, help='learning rate k-decay for cosine/poly (default: 1.0)') -parser.add_argument('--warmup-lr', type=float, default=0.0001, metavar='LR', +group.add_argument('--warmup-lr', type=float, default=0.0001, metavar='LR', help='warmup learning rate (default: 0.0001)') -parser.add_argument('--min-lr', type=float, default=1e-6, metavar='LR', +group.add_argument('--min-lr', type=float, default=1e-6, metavar='LR', help='lower lr bound for cyclic schedulers that hit 0 (1e-5)') -parser.add_argument('--epochs', type=int, default=300, metavar='N', +group.add_argument('--epochs', type=int, default=300, metavar='N', help='number of epochs to train (default: 300)') -parser.add_argument('--epoch-repeats', type=float, default=0., metavar='N', +group.add_argument('--epoch-repeats', type=float, default=0., metavar='N', help='epoch repeat multiplier (number of times to repeat dataset epoch per train epoch).') -parser.add_argument('--start-epoch', default=None, type=int, metavar='N', +group.add_argument('--start-epoch', default=None, type=int, metavar='N', help='manual epoch number (useful on restarts)') -parser.add_argument('--decay-milestones', default=[30, 60], type=int, nargs='+', metavar="MILESTONES", +group.add_argument('--decay-milestones', default=[30, 60], type=int, nargs='+', metavar="MILESTONES", help='list of decay epoch indices for multistep lr. must be increasing') -parser.add_argument('--decay-epochs', type=float, default=100, metavar='N', +group.add_argument('--decay-epochs', type=float, default=100, metavar='N', help='epoch interval to decay LR') -parser.add_argument('--warmup-epochs', type=int, default=3, metavar='N', +group.add_argument('--warmup-epochs', type=int, default=3, metavar='N', help='epochs to warmup LR, if scheduler supports') -parser.add_argument('--cooldown-epochs', type=int, default=10, metavar='N', +group.add_argument('--cooldown-epochs', type=int, default=10, metavar='N', help='epochs to cooldown LR at min_lr, after cyclic schedule ends') -parser.add_argument('--patience-epochs', type=int, default=10, metavar='N', +group.add_argument('--patience-epochs', type=int, default=10, metavar='N', help='patience epochs for Plateau LR scheduler (default: 10') -parser.add_argument('--decay-rate', '--dr', type=float, default=0.1, metavar='RATE', +group.add_argument('--decay-rate', '--dr', type=float, default=0.1, metavar='RATE', help='LR decay rate (default: 0.1)') # Augmentation & regularization parameters -parser.add_argument('--no-aug', action='store_true', default=False, +group = parser.add_argument_group('Augmentation and regularization parameters') +group.add_argument('--no-aug', action='store_true', default=False, help='Disable all training augmentation, override other train aug args') -parser.add_argument('--scale', type=float, nargs='+', default=[0.08, 1.0], metavar='PCT', +group.add_argument('--scale', type=float, nargs='+', default=[0.08, 1.0], metavar='PCT', help='Random resize scale (default: 0.08 1.0)') -parser.add_argument('--ratio', type=float, nargs='+', default=[3./4., 4./3.], metavar='RATIO', +group.add_argument('--ratio', type=float, nargs='+', default=[3./4., 4./3.], metavar='RATIO', help='Random resize aspect ratio (default: 0.75 1.33)') -parser.add_argument('--hflip', type=float, default=0.5, +group.add_argument('--hflip', type=float, default=0.5, help='Horizontal flip training aug probability') -parser.add_argument('--vflip', type=float, default=0., +group.add_argument('--vflip', type=float, default=0., help='Vertical flip training aug probability') -parser.add_argument('--color-jitter', type=float, default=0.4, metavar='PCT', +group.add_argument('--color-jitter', type=float, default=0.4, metavar='PCT', help='Color jitter factor (default: 0.4)') -parser.add_argument('--aa', type=str, default=None, metavar='NAME', +group.add_argument('--aa', type=str, default=None, metavar='NAME', help='Use AutoAugment policy. "v0" or "original". (default: None)'), -parser.add_argument('--aug-repeats', type=float, default=0, +group.add_argument('--aug-repeats', type=float, default=0, help='Number of augmentation repetitions (distributed training only) (default: 0)') -parser.add_argument('--aug-splits', type=int, default=0, +group.add_argument('--aug-splits', type=int, default=0, help='Number of augmentation splits (default: 0, valid: 0 or >=2)') -parser.add_argument('--jsd-loss', action='store_true', default=False, +group.add_argument('--jsd-loss', action='store_true', default=False, help='Enable Jensen-Shannon Divergence + CE loss. Use with `--aug-splits`.') -parser.add_argument('--bce-loss', action='store_true', default=False, +group.add_argument('--bce-loss', action='store_true', default=False, help='Enable BCE loss w/ Mixup/CutMix use.') -parser.add_argument('--bce-target-thresh', type=float, default=None, +group.add_argument('--bce-target-thresh', type=float, default=None, help='Threshold for binarizing softened BCE targets (default: None, disabled)') -parser.add_argument('--reprob', type=float, default=0., metavar='PCT', +group.add_argument('--reprob', type=float, default=0., metavar='PCT', help='Random erase prob (default: 0.)') -parser.add_argument('--remode', type=str, default='pixel', +group.add_argument('--remode', type=str, default='pixel', help='Random erase mode (default: "pixel")') -parser.add_argument('--recount', type=int, default=1, +group.add_argument('--recount', type=int, default=1, help='Random erase count (default: 1)') -parser.add_argument('--resplit', action='store_true', default=False, +group.add_argument('--resplit', action='store_true', default=False, help='Do not random erase first (clean) augmentation split') -parser.add_argument('--mixup', type=float, default=0.0, +group.add_argument('--mixup', type=float, default=0.0, help='mixup alpha, mixup enabled if > 0. (default: 0.)') -parser.add_argument('--cutmix', type=float, default=0.0, +group.add_argument('--cutmix', type=float, default=0.0, help='cutmix alpha, cutmix enabled if > 0. (default: 0.)') -parser.add_argument('--cutmix-minmax', type=float, nargs='+', default=None, +group.add_argument('--cutmix-minmax', type=float, nargs='+', default=None, help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)') -parser.add_argument('--mixup-prob', type=float, default=1.0, +group.add_argument('--mixup-prob', type=float, default=1.0, help='Probability of performing mixup or cutmix when either/both is enabled') -parser.add_argument('--mixup-switch-prob', type=float, default=0.5, +group.add_argument('--mixup-switch-prob', type=float, default=0.5, help='Probability of switching to cutmix when both mixup and cutmix enabled') -parser.add_argument('--mixup-mode', type=str, default='batch', +group.add_argument('--mixup-mode', type=str, default='batch', help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"') -parser.add_argument('--mixup-off-epoch', default=0, type=int, metavar='N', +group.add_argument('--mixup-off-epoch', default=0, type=int, metavar='N', help='Turn off mixup after this epoch, disabled if 0 (default: 0)') -parser.add_argument('--smoothing', type=float, default=0.1, +group.add_argument('--smoothing', type=float, default=0.1, help='Label smoothing (default: 0.1)') -parser.add_argument('--train-interpolation', type=str, default='random', +group.add_argument('--train-interpolation', type=str, default='random', help='Training interpolation (random, bilinear, bicubic default: "random")') -parser.add_argument('--drop', type=float, default=0.0, metavar='PCT', +group.add_argument('--drop', type=float, default=0.0, metavar='PCT', help='Dropout rate (default: 0.)') -parser.add_argument('--drop-connect', type=float, default=None, metavar='PCT', +group.add_argument('--drop-connect', type=float, default=None, metavar='PCT', help='Drop connect rate, DEPRECATED, use drop-path (default: None)') -parser.add_argument('--drop-path', type=float, default=None, metavar='PCT', +group.add_argument('--drop-path', type=float, default=None, metavar='PCT', help='Drop path rate (default: None)') -parser.add_argument('--drop-block', type=float, default=None, metavar='PCT', +group.add_argument('--drop-block', type=float, default=None, metavar='PCT', help='Drop block rate (default: None)') # Batch norm parameters (only works with gen_efficientnet based models currently) -parser.add_argument('--bn-momentum', type=float, default=None, +group = parser.add_argument_group('Batch norm parameters', 'Only works with gen_efficientnet based models currently.') +group.add_argument('--bn-momentum', type=float, default=None, help='BatchNorm momentum override (if not None)') -parser.add_argument('--bn-eps', type=float, default=None, +group.add_argument('--bn-eps', type=float, default=None, help='BatchNorm epsilon override (if not None)') -parser.add_argument('--sync-bn', action='store_true', +group.add_argument('--sync-bn', action='store_true', help='Enable NVIDIA Apex or Torch synchronized BatchNorm.') -parser.add_argument('--dist-bn', type=str, default='reduce', +group.add_argument('--dist-bn', type=str, default='reduce', help='Distribute BatchNorm stats between nodes after each epoch ("broadcast", "reduce", or "")') -parser.add_argument('--split-bn', action='store_true', +group.add_argument('--split-bn', action='store_true', help='Enable separate BN layers per augmentation split.') # Model Exponential Moving Average -parser.add_argument('--model-ema', action='store_true', default=False, +group = parser.add_argument_group('Model exponential moving average parameters') +group.add_argument('--model-ema', action='store_true', default=False, help='Enable tracking moving average of model weights') -parser.add_argument('--model-ema-force-cpu', action='store_true', default=False, +group.add_argument('--model-ema-force-cpu', action='store_true', default=False, help='Force ema to be tracked on CPU, rank=0 node only. Disables EMA validation.') -parser.add_argument('--model-ema-decay', type=float, default=0.9998, +group.add_argument('--model-ema-decay', type=float, default=0.9998, help='decay factor for model weights moving average (default: 0.9998)') # Misc -parser.add_argument('--seed', type=int, default=42, metavar='S', +group = parser.add_argument_group('Miscellaneous parameters') +group.add_argument('--seed', type=int, default=42, metavar='S', help='random seed (default: 42)') -parser.add_argument('--worker-seeding', type=str, default='all', +group.add_argument('--worker-seeding', type=str, default='all', help='worker seed mode (default: all)') -parser.add_argument('--log-interval', type=int, default=50, metavar='N', +group.add_argument('--log-interval', type=int, default=50, metavar='N', help='how many batches to wait before logging training status') -parser.add_argument('--recovery-interval', type=int, default=0, metavar='N', +group.add_argument('--recovery-interval', type=int, default=0, metavar='N', help='how many batches to wait before writing recovery checkpoint') -parser.add_argument('--checkpoint-hist', type=int, default=10, metavar='N', +group.add_argument('--checkpoint-hist', type=int, default=10, metavar='N', help='number of checkpoints to keep (default: 10)') -parser.add_argument('-j', '--workers', type=int, default=4, metavar='N', +group.add_argument('-j', '--workers', type=int, default=4, metavar='N', help='how many training processes to use (default: 4)') -parser.add_argument('--save-images', action='store_true', default=False, +group.add_argument('--save-images', action='store_true', default=False, help='save images of input bathes every log interval for debugging') -parser.add_argument('--amp', action='store_true', default=False, +group.add_argument('--amp', action='store_true', default=False, help='use NVIDIA Apex AMP or Native AMP for mixed precision training') -parser.add_argument('--apex-amp', action='store_true', default=False, +group.add_argument('--apex-amp', action='store_true', default=False, help='Use NVIDIA Apex AMP mixed precision') -parser.add_argument('--native-amp', action='store_true', default=False, +group.add_argument('--native-amp', action='store_true', default=False, help='Use Native Torch AMP mixed precision') -parser.add_argument('--no-ddp-bb', action='store_true', default=False, +group.add_argument('--no-ddp-bb', action='store_true', default=False, help='Force broadcast buffers for native DDP to off.') -parser.add_argument('--pin-mem', action='store_true', default=False, +group.add_argument('--pin-mem', action='store_true', default=False, help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.') -parser.add_argument('--no-prefetcher', action='store_true', default=False, +group.add_argument('--no-prefetcher', action='store_true', default=False, help='disable fast prefetcher') -parser.add_argument('--output', default='', type=str, metavar='PATH', +group.add_argument('--output', default='', type=str, metavar='PATH', help='path to output folder (default: none, current dir)') -parser.add_argument('--experiment', default='', type=str, metavar='NAME', +group.add_argument('--experiment', default='', type=str, metavar='NAME', help='name of train experiment, name of sub-folder for output') -parser.add_argument('--eval-metric', default='top1', type=str, metavar='EVAL_METRIC', +group.add_argument('--eval-metric', default='top1', type=str, metavar='EVAL_METRIC', help='Best metric (default: "top1"') -parser.add_argument('--tta', type=int, default=0, metavar='N', +group.add_argument('--tta', type=int, default=0, metavar='N', help='Test/inference time augmentation (oversampling) factor. 0=None (default: 0)') -parser.add_argument("--local_rank", default=0, type=int) -parser.add_argument('--use-multi-epochs-loader', action='store_true', default=False, +group.add_argument("--local_rank", default=0, type=int) +group.add_argument('--use-multi-epochs-loader', action='store_true', default=False, help='use the multi-epochs-loader to save time at the beginning of every epoch') -parser.add_argument('--log-wandb', action='store_true', default=False, +group.add_argument('--log-wandb', action='store_true', default=False, help='log training and validation metrics to wandb')