Add training utilities

__init__.py

@@ -1,3 +1,4 @@
 from .models.inception_resnet_v1 import InceptionResnetV1
 from .models.mtcnn import MTCNN, PNet, RNet, ONet, prewhiten
 from .models.utils.detect_face import extract_face
+from .models.utils import training

models/inception_resnet_v1.py

@@ -191,21 +191,24 @@ class InceptionResnetV1(nn.Module):
             (default: {None})
         classify {bool} -- Whether the model should output classification probabilities or feature
             embeddings. (default: {False})
-        num_classes {int} -- Number of output classes. Ignored if 'pretrained' is set, in which
-            case the number of classes is set to that used for training. (default: {1001})
+        num_classes {int} -- Number of output classes. If 'pretrained' is set and num_classes not
+            equal to that used for the pretrained model, the final linear layer will be randomly
+            initialized. (default: {1001})
+        dropout_prob {float} -- Dropout probability. (default: {0.6})
     """
-    def __init__(self, pretrained=None, classify=False, num_classes=1001):
+    def __init__(self, pretrained=None, classify=False, num_classes=1001, dropout_prob=0.6):
         super().__init__()
 
         # Set simple attributes
         self.pretrained = pretrained
         self.classify = classify
         self.num_classes = num_classes
 
+        tmp_classes = self.num_classes
         if pretrained == 'vggface2':
-            self.num_classes = 8631
+            tmp_classes = 8631
         elif pretrained == 'casia-webface':
-            self.num_classes = 10575
+            tmp_classes = 10575
 
         # Define layers
         self.conv2d_1a = BasicConv2d(3, 32, kernel_size=3, stride=2)
@@ -245,14 +248,16 @@ def __init__(self, pretrained=None, classify=False, num_classes=1001):
         )
         self.block8 = Block8(noReLU=True)
         self.avgpool_1a = nn.AdaptiveAvgPool2d(1)
+        self.dropout = nn.Dropout(dropout_prob)
         self.last_linear = nn.Linear(1792, 512, bias=False)
         self.last_bn = nn.BatchNorm1d(512, eps=0.001, momentum=0.1, affine=True)
-
-        self.logits = nn.Linear(512, self.num_classes)
-        self.softmax = nn.Softmax(dim=1)
+        self.logits = nn.Linear(512, tmp_classes)
 
         if pretrained is not None:
             load_weights(self, pretrained)
+
+        if self.num_classes != tmp_classes:
+            self.logits = nn.Linear(512, self.num_classes)
 
     def forward(self, x):
         """Calculate embeddings or probabilities given a batch of input image tensors.
@@ -277,12 +282,12 @@ def forward(self, x):
         x = self.repeat_3(x)
         x = self.block8(x)
         x = self.avgpool_1a(x)
+        x = self.dropout(x)
         x = self.last_linear(x.view(x.shape[0], -1))
         x = self.last_bn(x)
         x = F.normalize(x, p=2, dim=1)
         if self.classify:
             x = self.logits(x)
-            x = self.softmax(x)
         return x
 
 

models/utils/training.py

@@ -0,0 +1,136 @@
+import torch
+import numpy as np
+import time
+
+
+class Logger(object):
+
+    def __init__(self, mode, length, calculate_mean=False):
+        self.mode = mode
+        self.length = length
+        self.calculate_mean = calculate_mean
+        if self.calculate_mean:
+            self.fn = lambda x, i: x / (i + 1)
+        else:
+            self.fn = lambda x, i: x
+
+    def __call__(self, loss, metrics, i):
+        track_str = f'\r{self.mode} | {i + 1:5d}/{self.length:<5d}| '
+        loss_str = f'loss: {self.fn(loss, i):9.4f} | '
+        metric_str = ' | '.join(f'{k}: {self.fn(v, i):9.4f}' for k, v in metrics.items())
+        print(track_str + loss_str + metric_str + '   ', end='')
+        if i + 1 == self.length:
+            print('')
+
+
+class BatchTimer(object):
+    """Batch timing class.
+    Use this class for tracking training and testing time/rate per batch or per sample.
+    
+    Keyword Arguments:
+        rate {bool} -- Whether to report a rate (batches or samples per second) or a time (seconds
+            per batch or sample). (default: {True})
+        per_sample {bool} -- Whether to report times or rates per sample or per batch.
+            (default: {True})
+    """
+
+    def __init__(self, rate=True, per_sample=True):
+        self.start = time.time()
+        self.end = None
+        self.rate = rate
+        self.per_sample = per_sample
+
+    def __call__(self, y_pred, y):
+        self.end = time.time()
+        elapsed = self.end - self.start
+        self.start = self.end
+        self.end = None
+
+        if self.per_sample:
+            elapsed /= len(y_pred)
+        if self.rate:
+            elapsed = 1 / elapsed
+
+        return torch.tensor(elapsed)
+
+
+def accuracy(logits, y):
+    _, preds = torch.max(logits, 1)
+    return (preds == y).float().mean()
+
+
+def pass_epoch(
+    model, loss_fn, loader, optimizer=None, scheduler=None,
+    batch_metrics={'time': BatchTimer()}, show_running=True,
+    device='cpu', writer=None
+):
+    """Train or evaluate over a data epoch.
+    
+    Arguments:
+        model {torch.nn.Module} -- Pytorch model.
+        loss_fn {callable} -- A function to compute (scalar) loss.
+        loader {torch.utils.data.DataLoader} -- A pytorch data loader.
+    
+    Keyword Arguments:
+        optimizer {torch.optim.Optimizer} -- A pytorch optimizer.
+        scheduler {torch.optim.lr_scheduler._LRScheduler} -- LR scheduler (default: {None})
+        batch_metrics {dict} -- Dictionary of metric functions to call on each batch. The default
+            is a simple timer. A progressive average of these metrics, along with the average
+            loss, is printed every batch. (default: {{'time': iter_timer()}})
+        show_running {bool} -- Whether or not to print losses and metrics for the current batch
+            or rolling averages. (default: {False})
+        device {str or torch.device} -- Device for pytorch to use. (default: {'cpu'})
+        writer {torch.utils.tensorboard.SummaryWriter} -- Tensorboard SummaryWriter. (default: {None})
+    
+    Returns:
+        tuple(torch.Tensor, dict) -- A tuple of the average loss and a dictionary of average
+            metric values across the epoch.
+    """
+
+    mode = 'Train' if model.training else 'Valid'
+    logger = Logger(mode, length=len(loader), calculate_mean=show_running)
+    loss = 0
+    metrics = {}
+
+    for i_batch, (x, y) in enumerate(loader):
+        x = x.to(device)
+        y = y.to(device)
+        y_pred = model(x)
+        loss_batch = loss_fn(y_pred, y)
+
+        if model.training:
+            loss_batch.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+        metrics_batch = {}
+        for metric_name, metric_fn in batch_metrics.items():
+            metrics_batch[metric_name] = metric_fn(y_pred, y).detach().cpu()
+            metrics[metric_name] = metrics.get(metric_name, 0) + metrics_batch[metric_name]
+
+        if writer is not None and model.training:
+            if writer.iteration % writer.interval == 0:
+                writer.add_scalars('loss', {mode: loss_batch.detach().cpu()}, writer.iteration)
+                for metric_name, metric_batch in metrics_batch.items():
+                    writer.add_scalars(metric_name, {mode: metric_batch}, writer.iteration)
+            writer.iteration += 1
+
+        loss_batch = loss_batch.detach().cpu()
+        loss += loss_batch
+        if show_running:
+            logger(loss, metrics, i_batch)
+        else:
+            logger(loss_batch, metrics_batch, i_batch)
+
+    if model.training and scheduler is not None:
+        scheduler.step()
+
+    loss = loss / (i_batch + 1)
+    metrics = {k: v / (i_batch + 1) for k, v in metrics.items()}
+
+    if writer is not None and not model.training:
+        writer.add_scalars('loss', {mode: loss.detach()}, writer.iteration)
+        for metric_name, metric in metrics.items():
+            writer.add_scalars(metric_name, {mode: metric})
+
+    return loss, metrics

setup.py

@@ -1,23 +1,34 @@
-import setuptools
+import setuptools, os
 
-with open('facenet_pytorch/README.md', 'r') as f:
+PACKAGE_NAME = 'facenet-pytorch'
+VERSION = '0.2.2'
+AUTHOR = 'Tim Esler'
+EMAIL = '[email protected]'
+DESCRIPTION = 'Pretrained Pytorch face detection and recognition models'
+GITHUB_URL = 'https://github.com/timesler/facenet-pytorch'
+
+parent_dir = os.path.dirname(os.path.realpath(__file__))
+import_name = os.path.basename(parent_dir)
+
+with open(f'{parent_dir}/README.md', 'r') as f:
     long_description = f.read()
 
 setuptools.setup(
-    name='facenet-pytorch',
-    version='0.1.0',
-    author='Tim Esler',
-    author_email='[email protected]',
-    description='Pretrained Pytorch face detection and recognition models',
+    name=PACKAGE_NAME,
+    version=VERSION,
+    author=AUTHOR,
+    author_email=EMAIL,
+    description=DESCRIPTION,
     long_description=long_description,
     long_description_content_type='text/markdown',
-    url='https://github.com/timesler/facenet-pytorch',
+    url=GITHUB_URL,
     packages=[
         'facenet_pytorch',
         'facenet_pytorch.models',
         'facenet_pytorch.models.utils',
         'facenet_pytorch.data',
-   ],
+    ],
+    package_dir={'facenet_pytorch':'.'},
     package_data={'': ['*net.pt']},
     classifiers=[
         "Programming Language :: Python :: 3",

tests/travis_test.py

@@ -118,8 +118,6 @@ def get_image(path, trans):
 
 resnet_pt = InceptionResnetV1(pretrained=ds, classify=True).eval()
 prob = resnet_pt(aligned)
-if sys.platform != 'win32':
-    assert prob.mean().detach().item() - 9.4563e-05 < 1e-5
 
 # MULTI-FACE TEST