Cleanup (#11)

* cifar wrapper

* move load_state_dict to nbdt.models.utils

* train loss oops

* cleaner dataset kwarg attempt

* aesthetic changes

* decorator for analyzer

* imports and docstring

* moar aesthetics

* more aesthetic

main.py

@@ -1,50 +1,33 @@
 """
-Neural-Backed Decision Trees training script on CIFAR10, CIFAR100, TinyImagenet200
-
-The original version of this `main.py` was taken from
-
- https://github.com/kuangliu/pytorch-cifar
-
-and extended in
-
- https://github.com/alvinwan/pytorch-cifar-plus
+Neural-Backed Decision Trees training on CIFAR10, CIFAR100, TinyImagenet200
 
+The original version of this `main.py` was taken from kuangliu/pytorch-cifar.
 The script has since been heavily modified to support a number of different
-configurations and options. See the current repository for a full description
-of its bells and whistles.
-
- https://github.com/alvinwan/neural-backed-decision-trees
+configurations and options: alvinwan/neural-backed-decision-trees
 """
+import os
+import argparse
+import numpy as np
 import torch
-import torch.nn as nn
-import torch.optim as optim
+from torch import nn, optim
 import torch.nn.functional as F
 import torch.backends.cudnn as cudnn
-from nbdt import data, analysis, loss, models, metrics
-
 import torchvision
 import torchvision.transforms as transforms
 
-import os
-import argparse
-import numpy as np
-
-from nbdt.utils import (
- progress_bar, generate_checkpoint_fname, generate_kwargs, Colors
-)
+from nbdt import data, analysis, loss, models, metrics
+from nbdt.utils import progress_bar, generate_checkpoint_fname, generate_kwargs, Colors
 from nbdt.thirdparty.wn import maybe_install_wordnet
+from nbdt.models.utils import load_state_dict, make_kwarg_optional
 
 maybe_install_wordnet()
 
-datasets = ('CIFAR10', 'CIFAR100') + data.imagenet.names + data.custom.names
-
-
 parser = argparse.ArgumentParser(description='PyTorch CIFAR Training')
 parser.add_argument('--batch-size', default=512, type=int,
  help='Batch size used for training')
 parser.add_argument('--epochs', '-e', default=200, type=int,
  help='By default, lr schedule is scaled accordingly')
-parser.add_argument('--dataset', default='CIFAR10', choices=datasets)
+parser.add_argument('--dataset', default='CIFAR10', choices=data.cifar.names + data.imagenet.names + data.custom.names)
 parser.add_argument('--arch', default='ResNet18', choices=list(models.get_model_choices()))
 parser.add_argument('--lr', default=0.1, type=float, help='learning rate')
 parser.add_argument('--resume', '-r', action='store_true', help='resume from checkpoint')
@@ -57,22 +40,16 @@
 parser.add_argument('--pretrained', action='store_true',
  help='Download pretrained model. Not all models support this.')
 parser.add_argument('--eval', help='eval only', action='store_true')
-
-# options specific to this project and its dataloaders
 parser.add_argument('--loss', choices=loss.names, default='CrossEntropyLoss')
 parser.add_argument('--metric', choices=metrics.names, default='top1')
 parser.add_argument('--analysis', choices=analysis.names, help='Run analysis after each epoch')
-parser.add_argument('--input-size', type=int,
- help='Set transform train and val. Samples are resized to '
- 'input-size + 32.')
-parser.add_argument('--lr-decay-every', type=int, default=0)
 
+# other dataset, loss or analysis specific options
 data.custom.add_arguments(parser)
 loss.add_arguments(parser)
 analysis.add_arguments(parser)
 
 args = parser.parse_args()
-
 loss.set_default_values(args)
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -81,40 +58,9 @@
 
 # Data
 print('==> Preparing data..')
-transform_train = transforms.Compose([
- transforms.RandomCrop(32, padding=4),
- transforms.RandomHorizontalFlip(),
- transforms.ToTensor(),
- transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
-])
-
-transform_test = transforms.Compose([
- transforms.ToTensor(),
- transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
-])
-
 dataset = getattr(data, args.dataset)
-
-if args.dataset in ('TinyImagenet200', 'Imagenet1000'):
- default_input_size = 64 if args.dataset == 'TinyImagenet200' else 224
- input_size = args.input_size or default_input_size
- transform_train = dataset.transform_train(input_size)
- transform_test = dataset.transform_val(input_size)
-elif args.input_size is not None and args.input_size > 32:
- transform_train = transforms.Compose([
- transforms.Resize(args.input_size + 32),
- transforms.RandomCrop(args.input_size),
- transforms.RandomHorizontalFlip(),
- transforms.ToTensor(),
- transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
- ])
-
- transform_test = transforms.Compose([
- transforms.Resize(args.input_size + 32),
- transforms.CenterCrop(args.input_size),
- transforms.ToTensor(),
- transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
- ])
+transform_train = dataset.transform_train()
+transform_test = dataset.transform_val()
 
 dataset_kwargs = generate_kwargs(args, dataset,
  name=f'Dataset {args.dataset}',
@@ -138,16 +84,7 @@
 
 if args.pretrained:
  print('==> Loading pretrained model..')
- try:
- net = model(pretrained=True, dataset=args.dataset, **model_kwargs)
- except TypeError as e: # likely because `dataset` not allowed arg
- print(e)
-
- try:
- net = model(pretrained=True, **model_kwargs)
- except Exception as e:
- Colors.red(f'Fatal error: {e}')
- exit()
+ net = make_kwarg_optional(model, dataset=args.dataset)(pretrained=True, **model_kwargs)
 else:
  net = model(**model_kwargs)
 
@@ -160,19 +97,6 @@
 checkpoint_path = './checkpoint/{}.pth'.format(checkpoint_fname)
 print(f'==> Checkpoints will be saved to: {checkpoint_path}')
 
-
-# TODO(alvin): fix checkpoint structure so that this isn't neededd
-def load_state_dict(state_dict):
- try:
- net.load_state_dict(state_dict)
- except RuntimeError as e:
- if 'Missing key(s) in state_dict:' in str(e):
- net.load_state_dict({
- key.replace('module.', '', 1): value
- for key, value in state_dict.items()
- })
-
-
 resume_path = args.path_resume or checkpoint_path
 if args.resume:
  # Load checkpoint.
@@ -184,13 +108,13 @@ def load_state_dict(state_dict):
  checkpoint = torch.load(resume_path, map_location=torch.device(device))
 
  if 'net' in checkpoint:
- load_state_dict(checkpoint['net'])
+ load_state_dict(net, checkpoint['net'])
  best_acc = checkpoint['acc']
  start_epoch = checkpoint['epoch']
  Colors.cyan(f'==> Checkpoint found for epoch {start_epoch} with accuracy '
  f'{best_acc} at {resume_path}')
  else:
- load_state_dict(checkpoint)
+ load_state_dict(net, checkpoint)
  Colors.cyan(f'==> Checkpoint found at {resume_path}')
 
 
@@ -205,19 +129,25 @@ def load_state_dict(state_dict):
 optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
 
 def adjust_learning_rate(epoch, lr):
- if args.lr_decay_every:
- steps = epoch // args.lr_decay_every
- return lr / (10 ** steps)
  if epoch <= 150 / 350. * args.epochs: # 32k iterations
  return lr
  elif epoch <= 250 / 350. * args.epochs: # 48k iterations
  return lr/10
  else:
  return lr/100
 
+
+class_analysis = getattr(analysis, args.analysis or 'Noop')
+analyzer_kwargs = generate_kwargs(args, class_analysis,
+ name=f'Analyzer {args.analysis}',
+ keys=analysis.keys,
+ globals=globals())
+analyzer = class_analysis(**analyzer_kwargs)
+
+
 # Training
+@analyzer.train_function
 def train(epoch, analyzer):
- analyzer.start_train(epoch)
  lr = adjust_learning_rate(epoch, args.lr)
  optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
 
@@ -235,19 +165,14 @@ def train(epoch, analyzer):
 
  train_loss += loss.item()
  metric.forward(outputs, targets)
-
  stat = analyzer.update_batch(outputs, targets)
- extra = f'| {stat}' if stat else ''
 
- progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d) %s'
- % (test_loss/(batch_idx+1), 100. * metric.report(),
- metric.correct, metric.total, extra))
+ progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d) %s' % (
+ train_loss / ( batch_idx + 1 ), 100. * metric.report(), metric.correct, metric.total, f'| {stat}' if stat else ''))
 
- analyzer.end_train(epoch)
 
+@analyzer.test_function
 def test(epoch, analyzer, checkpoint=True):
- analyzer.start_test(epoch)
-
  global best_acc
  net.eval()
  test_loss = 0
@@ -260,17 +185,10 @@ def test(epoch, analyzer, checkpoint=True):
 
  test_loss += loss.item()
  metric.forward(outputs, targets)
-
- if device == 'cuda':
- outputs = outputs.cpu()
- targets = targets.cpu()
-
  stat = analyzer.update_batch(outputs, targets)
- extra = f'| {stat}' if stat else ''
 
- progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d) %s'
- % (test_loss/(batch_idx+1), 100. * metric.report(),
- metric.correct, metric.total, extra))
+ progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d) %s' % (
+ test_loss / ( batch_idx + 1 ), 100. * metric.report(), metric.correct, metric.total, f'| {stat}' if stat else ''))
 
  # Save checkpoint.
  acc = 100. * metric.report()
@@ -288,16 +206,6 @@ def test(epoch, analyzer, checkpoint=True):
  torch.save(state, f'./checkpoint/{checkpoint_fname}.pth')
  best_acc = acc
 
- analyzer.end_test(epoch)
-
-
-class_analysis = getattr(analysis, args.analysis or 'Noop')
-analyzer_kwargs = generate_kwargs(args, class_analysis,
- name=f'Analyzer {args.analysis}',
- keys=analysis.keys,
- globals=globals())
-analyzer = class_analysis(**analyzer_kwargs)
-
 
 if args.eval:
  if not args.resume and not args.pretrained:
@@ -306,16 +214,13 @@ def test(epoch, analyzer, checkpoint=True):
 
  analyzer.start_epoch(0)
  test(0, analyzer, checkpoint=False)
+ analyzer.end_epoch(0)
  exit()
 
+
 for epoch in range(start_epoch, args.epochs):
  analyzer.start_epoch(epoch)
  train(epoch, analyzer)
  test(epoch, analyzer)
  analyzer.end_epoch(epoch)
-
-if args.epochs == 0:
- analyzer.start_epoch(0)
- test(0, analyzer)
- analyzer.end_epoch(0)
 print(f'Best accuracy: {best_acc} // Checkpoint name: {checkpoint_fname}')

nbdt/analysis.py

@@ -4,6 +4,7 @@
  HardEmbeddedDecisionRules as HardRules
 )
 from nbdt import metrics
+import functools
 import numpy as np
 
 
@@ -17,6 +18,20 @@ def add_arguments(parser):
  pass
 
 
+def start_end_decorator(obj, name):
+ start = getattr(obj, f'start_{name}', None)
+ end = getattr(obj, f'end_{name}', None)
+ assert start and end
+ def decorator(f):
+ @functools.wraps(f)
+ def wrapper(epoch, *args, **kwargs):
+ start(epoch)
+ f(epoch, *args, **kwargs)
+ end(epoch)
+ return wrapper
+ return decorator
+
+
 class Noop:
 
  accepts_classes = lambda trainset, **kwargs: trainset.classes
@@ -28,6 +43,18 @@ def __init__(self, classes=()):
  self.num_classes = len(classes)
  self.epoch = None
 
+ @property
+ def epoch_function(self):
+ return start_end_decorator(self, 'epoch')
+
+ @property
+ def train_function(self):
+ return start_end_decorator(self, 'train')
+
+ @property
+ def test_function(self):
+ return start_end_decorator(self, 'test')
+
  def start_epoch(self, epoch):
  self.epoch = epoch
 

nbdt/data/__init__.py

@@ -4,3 +4,4 @@
 from .lip import *
 from .ade20k import *
 from torchvision.datasets import *
+from .cifar import *

nbdt/data/cifar.py

@@ -0,0 +1,31 @@
+"""Wrappers around CIFAR datasets"""
+
+from torchvision import datasets, transforms
+
+__all__ = names = ('CIFAR10', 'CIFAR100')
+
+
+class CIFAR:
+
+ @staticmethod
+ def transform_train():
+ return transforms.Compose([
+ transforms.RandomCrop(32, padding=4),
+ transforms.RandomHorizontalFlip(),
+ transforms.ToTensor(),
+ transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+ ])
+
+ @staticmethod
+ def transform_val():
+ return transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+ ])
+
+
+class CIFAR10(datasets.CIFAR10, CIFAR):
+ pass
+
+class CIFAR100(datasets.CIFAR100, CIFAR):
+ pass