Merge pull request #4 from Irynei/feature/enhance_training

Feature/enhance training

Author: Irynei

base/base_dataset.py

@@ -10,11 +10,12 @@ class AutoAugmentDataset(data.Dataset):
  Randomly applies subset of augmentations and set them as labels
 
  """
- def __init__(self, dataset, base_transforms, augmentations, max_size=7):
+ def __init__(self, dataset, base_transforms, augmentations, max_size=7, train=True):
  self.dataset = dataset
  self.base_transforms = base_transforms
  self.augmentations = augmentations
  self.max_size = max_size
+ self.train = train
 
  def __getitem__(self, index):
  x, y = self.dataset[index]
@@ -35,22 +36,30 @@ def __len__(self):
 
  def get_subset_of_transforms(self):
  """
- Randomly get size of subset and then randomly choose subset of transformations
+ in case of train dataset:
+ Randomly get size of subset and then randomly choose subset of transformations
+ in case of test dataset:
+ Subset of transformations is always empty
 
  Returns:
  list of chosen transformations, one-hot-encoded labels
 
  """
  all_transforms_size = len(self.augmentations)
 
- # size from 0 to max_size - 1
- subset_size = np.random.randint(0, self.max_size)
- all_transforms_idx = np.arange(all_transforms_size)
- # get random subset without duplicates
- np.random.shuffle(all_transforms_idx)
- transform_idx = all_transforms_idx[:subset_size]
- subset_transforms = [self.augmentations[i] for i in transform_idx]
+ if self.train:
+ # size from 0 to max_size - 1
+ subset_size = np.random.randint(0, self.max_size)
+ all_transforms_idx = np.arange(all_transforms_size)
+ # get random subset without duplicates
+ np.random.shuffle(all_transforms_idx)
+ transform_idx = all_transforms_idx[:subset_size]
+ subset_transforms = [self.augmentations[i] for i in transform_idx]
 
- labels = np.zeros(all_transforms_size)
- labels[transform_idx] = 1
+ labels = np.zeros(all_transforms_size)
+ labels[transform_idx] = 1
+ else:
+ # in case of test we do
+ labels = np.zeros(all_transforms_size)
+ subset_transforms = []
  return subset_transforms, labels

base/base_trainer.py

@@ -3,6 +3,7 @@
 import torch
 import glog as log
 import torch.optim as optim
+from utils.util import EarlyStopping
 
 
 class BaseTrainer:
@@ -50,6 +51,10 @@ def __init__(self, model, loss, metrics, resume, config, train_logger=None):
  assert self.monitor_mode == 'min' or self.monitor_mode == 'max'
  self.monitor_best = math.inf if self.monitor_mode == 'min' else -math.inf
 
+ self.early_stopping = None
+ if config.get('early_stopping'):
+ self.early_stopping = EarlyStopping(**config['early_stopping']['early_stopping_params'])
+
  self.start_epoch = 1
  self.checkpoint_dir = os.path.join(config['trainer']['save_dir'], config['experiment_name'])
  if resume:
@@ -93,9 +98,16 @@ def train(self):
 
  # lr_scheduler logic
  if self.lr_scheduler and epoch % self.lr_scheduler_freq == 0:
- self.lr_scheduler.step(epoch)
- lr = self.lr_scheduler.get_lr()[0]
- self.logger.info('New Learning Rate: {:.6f}'.format(lr))
+ if isinstance(self.lr_scheduler, optim.lr_scheduler.ReduceLROnPlateau):
+ self.lr_scheduler.step(log['val_loss'])
+ else:
+ self.lr_scheduler.step(epoch)
+
+ # stopping early logic
+ if self.early_stopping:
+ stop_early = self.early_stopping.step(log['val_loss'])
+ if stop_early:
+ break
 
  def _train_epoch(self, epoch):
  """

configs/densenet121_on_cifar10.json

@@ -0,0 +1,55 @@
+{
+ "experiment_name": "DenseNet121_on_CIFAR10",
+ "cuda": true,
+ "gpu": 0,
+ "augmentation": {
+ "max_size": 5
+ },
+ "data_loader": {
+ "name": "CIFAR10DataLoader",
+ "data_dir": "datasets/",
+ "batch_size": 32,
+ "shuffle": true
+ },
+ "model_name": "densenet121_32x32",
+ "model_params": {
+ "num_classes": 15
+ },
+ "optimizer_type": "SGD",
+ "optimizer_params": {
+ "lr": 0.1,
+ "weight_decay": 0.0005,
+ "momentum": 0.9
+ },
+ "loss": "MultiLabelSoftMarginLoss",
+ "validation": {
+ "validation_split": 0.1,
+ "shuffle": true
+ },
+ "lr_scheduler": {
+ "lr_scheduler_type": "ReduceLROnPlateau",
+ "lr_scheduler_freq": 1,
+ "additional_params": {
+ "patience": 8,
+ "mode": "min",
+ "min_lr": 1e-7,
+ "factor": 0.1,
+ "verbose": true
+ }
+ },
+ "early_stopping": {
+ "early_stopping_params": {
+ "patience": 12,
+ "mode": "min"
+ }
+ },
+ "metrics": ["accuracy", "jaccard_similarity"],
+ "trainer": {
+ "epochs": 500,
+ "save_dir": "experiments/",
+ "save_freq": 100,
+ "verbosity": 2,
+ "monitor": "val_loss",
+ "monitor_mode": "min"
+ }
+}

data_loaders/data_loader.py

@@ -42,13 +42,14 @@ def __init__(self, config):
  dataset=datasets.CIFAR10(self.data_dir, train=True, download=True),
  base_transforms=self.base_transforms,
  augmentations=self.augmentations,
- max_size=self.max_size
+ max_size=self.max_size,
+ train=True
  ),
- 'test': datasets.CIFAR10(
- self.data_dir,
- train=False,
- download=True,
- transform=self.base_transforms
+ 'test': AutoAugmentDataset(
+ dataset=datasets.CIFAR10(self.data_dir, train=False, download=True),
+ base_transforms=self.base_transforms,
+ augmentations=self.augmentations,
+ train=False
  )
  }
  super(CIFAR10DataLoader, self).__init__(self.dataset, config)
@@ -74,11 +75,11 @@ def __init__(self, config):
  augmentations=self.augmentations,
  max_size=self.max_size
  ),
- 'test': datasets.SVHN(
- self.data_dir,
- split='test',
- download=True,
- transform=self.base_transforms
+ 'test': AutoAugmentDataset(
+ dataset=datasets.SVHN(self.data_dir, split='test', download=True),
+ base_transforms=self.base_transforms,
+ augmentations=self.augmentations,
+ train=False
  )
  }
  super(SVHNDataLoader, self).__init__(self.dataset, config)

model/architectures/__init__.py

@@ -0,0 +1,9 @@
+from .densenet_32x32 import (
+ densenet121_32x32,
+ densenet161_32x32,
+ densenet169_32x32,
+ densenet201_32x32
+)
+from .vgg_32x32 import (
+ VGG16_32x32
+)

model/architectures/densenet_32x32.py

@@ -0,0 +1,102 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def densenet121_32x32(num_classes):
+ """ densenet121 that works with 32x32 input (e.g. CIFAR10) """
+ return DenseNet(Bottleneck, [6, 12, 24, 16], growth_rate=32, num_classes=num_classes)
+
+
+def densenet169_32x32(num_classes):
+ """ densenet169 that works with 32x32 input (e.g. CIFAR10) """
+ return DenseNet(Bottleneck, [6, 12, 32, 32], growth_rate=32, num_classes=num_classes)
+
+
+def densenet201_32x32(num_classes):
+ """ densenet201 that works with 32x32 input (e.g. CIFAR10) """
+ return DenseNet(Bottleneck, [6, 12, 48, 32], growth_rate=32, num_classes=num_classes)
+
+
+def densenet161_32x32(num_classes):
+ """ densenet161 that works with 32x32 input (e.g. CIFAR10) """
+ return DenseNet(Bottleneck, [6, 12, 36, 24], growth_rate=48, num_classes=num_classes)
+
+
+class Bottleneck(nn.Module):
+ def __init__(self, in_planes, growth_rate):
+ super(Bottleneck, self).__init__()
+ self.bn1 = nn.BatchNorm2d(in_planes)
+ self.conv1 = nn.Conv2d(in_planes, 4 * growth_rate, kernel_size=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(4 * growth_rate)
+ self.conv2 = nn.Conv2d(4 * growth_rate, growth_rate, kernel_size=3, padding=1, bias=False)
+
+ def forward(self, x):
+ out = self.conv1(F.relu(self.bn1(x)))
+ out = self.conv2(F.relu(self.bn2(out)))
+ out = torch.cat([out, x], 1)
+ return out
+
+
+class Transition(nn.Module):
+ def __init__(self, in_planes, out_planes):
+ super(Transition, self).__init__()
+ self.bn = nn.BatchNorm2d(in_planes)
+ self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=1, bias=False)
+
+ def forward(self, x):
+ out = self.conv(F.relu(self.bn(x)))
+ out = F.avg_pool2d(out, 2)
+ return out
+
+
+class DenseNet(nn.Module):
+ def __init__(self, block, nblocks, growth_rate=12, reduction=0.5, num_classes=10):
+ super(DenseNet, self).__init__()
+ self.growth_rate = growth_rate
+
+ num_planes = 2 * growth_rate
+ self.conv1 = nn.Conv2d(3, num_planes, kernel_size=3, padding=1, bias=False)
+
+ self.dense1 = self._make_dense_layers(block, num_planes, nblocks[0])
+ num_planes += nblocks[0] * growth_rate
+ out_planes = int(math.floor(num_planes * reduction))
+ self.trans1 = Transition(num_planes, out_planes)
+ num_planes = out_planes
+
+ self.dense2 = self._make_dense_layers(block, num_planes, nblocks[1])
+ num_planes += nblocks[1] * growth_rate
+ out_planes = int(math.floor(num_planes * reduction))
+ self.trans2 = Transition(num_planes, out_planes)
+ num_planes = out_planes
+
+ self.dense3 = self._make_dense_layers(block, num_planes, nblocks[2])
+ num_planes += nblocks[2] * growth_rate
+ out_planes = int(math.floor(num_planes * reduction))
+ self.trans3 = Transition(num_planes, out_planes)
+ num_planes = out_planes
+
+ self.dense4 = self._make_dense_layers(block, num_planes, nblocks[3])
+ num_planes += nblocks[3] * growth_rate
+
+ self.bn = nn.BatchNorm2d(num_planes)
+ self.linear = nn.Linear(num_planes, num_classes)
+
+ def _make_dense_layers(self, block, in_planes, nblock):
+ layers = []
+ for i in range(nblock):
+ layers.append(block(in_planes, self.growth_rate))
+ in_planes += self.growth_rate
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+ out = self.conv1(x)
+ out = self.trans1(self.dense1(out))
+ out = self.trans2(self.dense2(out))
+ out = self.trans3(self.dense3(out))
+ out = self.dense4(out)
+ out = F.avg_pool2d(F.relu(self.bn(out)), 4)
+ out = out.view(out.size(0), -1)
+ out = self.linear(out)
+ return out

model/architectures/vgg_32x32.py

@@ -0,0 +1,30 @@
+import torch.nn as nn
+
+
+class VGG16_32x32(nn.Module):
+ """ VGG16 that works with 32x32 input (e.g. CIFAR10)"""
+ def __init__(self, num_classes):
+ super(VGG16_32x32, self).__init__()
+ self.layers = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
+ self.features = self._make_layers(self.layers)
+ self.classifier = nn.Linear(512, num_classes)
+
+ def forward(self, x):
+ out = self.features(x)
+ out = out.view(out.size(0), -1)
+ out = self.classifier(out)
+ return out
+
+ def _make_layers(self, cfg):
+ layers = []
+ in_channels = 3
+ for x in cfg:
+ if x == 'M':
+ layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
+ else:
+ layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
+ nn.BatchNorm2d(x),
+ nn.ReLU(inplace=True)]
+ in_channels = x
+ layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
+ return nn.Sequential(*layers)

model/metric.py

@@ -3,6 +3,7 @@
  accuracy_score,
  fbeta_score,
  hamming_loss,
+ jaccard_similarity_score
 )
 
 
@@ -18,14 +19,16 @@ def get_metric_functions(metric_names):
  metric_fns.append(f_beta)
  elif metric_name == 'ham_loss':
  metric_fns.append(ham_loss)
+ elif metric_name == 'jaccard_similarity':
+ metric_fns.append(jaccard_similarity)
  else:
  raise NameError("Metric '{metric}' not found.".format(metric=metric_name))
  return metric_fns
 
 
 def accuracy(preds, targs, threshold=0.5):
  """
- Accuracy classification score.
+ Exact match accuracy classification score.
  The set of labels predicted for a sample (preds) must exactly match the
  corresponding set of labels (targs)
  Args:
@@ -64,3 +67,16 @@ def ham_loss(preds, targs, threshold=0.5):
 
  """
  return hamming_loss(targs, (preds > threshold))
+
+
+def jaccard_similarity(preds, targs, threshold=0.5):
+ """
+ Jaccard similarity score for multi-label classification.
+ Intersection over union
+ Args:
+ preds: predicted targets as returned by a model
+ targs: ground truth target value
+ threshold: threshold, default is 0.5
+
+ """
+ return jaccard_similarity_score(targs, (preds > threshold))