ResNext models

vectorizer/README.md

@@ -16,7 +16,7 @@ Set environment variables:
 | --- | --- | --- |
 | VS_PORT | 8080 | Port to listen (for Flask) |
 | VS_FAN_MODEL |  | Path to identification model |
-| VS_REC_DEPTH | 50 | Recognition net depth |
+| VS_REC_NET | resnet50 | Recognition net name |
 | VS_REC_MODEL |  | Path to recognition model |
 
 Do not change configuration if you want run prepared docker-compose.
@@ -112,7 +112,7 @@ python3 -m recognition.train \
 --lfw_root ~/datasets/lfw \
 --lfw_pair_list lfw_test_pair.txt \
 --model_name recongition1 --batch_size 20 \
---loss adacos --print_freq 20 --depth 50
+--loss adacos --print_freq 20 --net resnet50
 ```
 
 | Argument | Description | Required / Default value |
@@ -121,7 +121,7 @@ python3 -m recognition.train \
 | --casia_root | Path to CASIA images | Yes |
 | --lfw_root | Path to LFW dataset | Yes |
 | --lfw_pair_list | Path to LFW pair list file (lfw_test_pair.txt - in this repository) | Yes |
-| --depth | Resnet depth, must be one of 18, 34, 50, 101, 152 or 20 for sphere net | 50 |
+| --net | Net name, must be one of resnet18, resnet34, resnet50, resnet101, resnet152, resnext50, resnext101 or spherenet | resnet50 |
 | --epochs | Number of epochs | 50 |
 | --batch_size | Batch size | 16 |
 | --model_name | Model name prefix | Yes |

vectorizer/recognition/nets.py

@@ -74,21 +74,43 @@ def sphere20():
     return sphere20a()
 
 
-def get_net_by_depth(depth):
+def resnext50(pretrained=False, **kwargs):
-    if depth == 18:
+    """Constructs a ResNext-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = models.resnext50_32x4d(num_classes=512, **kwargs)
+    return model
+
+
+def resnext101(pretrained=False, **kwargs):
+    """Constructs a ResNext-101 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = models.resnext101_32x8d(num_classes=512, **kwargs)
+    return model
+
+
+def get_net_by_name(name):
+    if name == 'resnet18':
         model = resnet18()
-    elif depth == 20:
+    elif name == 'spherenet':
         model = sphere20()
-    elif depth == 34:
+    elif name == 'resnet34':
         model = resnet34()
-    elif depth == 50:
+    elif name == 'resnet50':
         model = resnet50()
-    elif depth == 101:
+    elif name == 'resnet101':
         model = resnet101()
-    elif depth == 152:
+    elif name == 'resnet152':
         model = resnet152()
+    elif name == 'resnext50':
+        model = resnext50()
+    elif name == 'resnext101':
+        model = resnext101()
     else:
-        raise ValueError('Unsupported model depth %d, must be one of 18, 34, 50, 101, 152' % depth)
+        raise ValueError('Unsupported model %s, must be one of resnet18, resnet34, resnet50, resnet101, resnet152, spherenet, resnext50, resnext101' % name)
 
     return model
 

vectorizer/recognition/test.py

@@ -25,7 +25,7 @@ import argparse
 
 from torch.utils.data import TensorDataset, DataLoader
 
-from recognition.nets import get_net_by_depth
+from recognition.nets import get_net_by_name
 import torch
 import numpy as np
 from torch.nn import DataParallel
@@ -135,8 +135,8 @@ def cal_accuracy(y_score, y_true):
 def main(args=None):
     parser = argparse.ArgumentParser(description='Testing script for face identification.')
 
-    parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152 or 20 for sphere', type=int,
+    parser.add_argument('--net', help='Net name, must be one of resnet18, resnet34, resnet50 (default), resnet101, resnet152, resnext50, resnext101 or spherenet',
-                        default=50)
+                        default='resnet50')
     parser.add_argument('--parallel', help='Run training with DataParallel', dest='parallel',
                         default=False, action='store_true')
     parser.add_argument('--model', help='Path to model')
@@ -149,7 +149,7 @@ def main(args=None):
     is_cuda = torch.cuda.is_available()
     print('CUDA available: {}'.format(is_cuda))
 
-    model = get_net_by_depth(parser.depth)
+    model = get_net_by_name(parser.net)
 
     if parser.parallel:
         model = DataParallel(model)

vectorizer/recognition/train.py

@@ -30,7 +30,7 @@ from torchvision import transforms as T
 
 from recognition.angle import AngleLinear, CosFace, SphereFace, ArcFace, AdaCos
 from recognition.focal_loss import FocalLoss
-from recognition.nets import get_net_by_depth
+from recognition.nets import get_net_by_name
 from recognition.test import lfw_test2, get_pair_list, load_img_data
 
 
@@ -81,8 +81,8 @@ def main(args=None):
 
     parser.add_argument('--print_freq', help='Print every N batch (default 100)', type=int, default=100)
     parser.add_argument('--epochs', help='Number of epochs', type=int, default=50)
-    parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152 or 20 for sphere', type=int,
+    parser.add_argument('--net', help='Net name, must be one of resnet18, resnet34, resnet50, resnet101, resnet152, resnext50, resnext101 or spherenet',
-                        default=50)
+                        default='resnet50')
     parser.add_argument('--lr_step', help='Learning rate step (default 10)', type=int, default=10)
     parser.add_argument('--lr', help='Learning rate (default 0.1)', type=float, default=0.1)
     parser.add_argument('--weight_decay', help='Weight decay (default 0.0005)', type=float, default=0.0005)
@@ -108,7 +108,7 @@ def main(args=None):
     print('CUDA available: {}'.format(is_cuda))
 
     imagesize = 224
-    model = get_net_by_depth(parser.depth)
+    model = get_net_by_name(parser.net)
 
     # TODO split training dataset to train/validation and stop using test dataset for acc
     train_dataset = Dataset(parser.casia_root, parser.casia_list, imagesize)
@@ -172,8 +172,6 @@ def main(args=None):
     start = time.time()
     last_acc = 0.0
     for i in range(parser.epochs):
-        scheduler.step()
-
         model.train()
         for ii, data in enumerate(trainloader):
             data_input, label = data
@@ -196,6 +194,7 @@ def main(args=None):
 
                 start = time.time()
 
+        scheduler.step()
         model.eval()
         acc = lfw_test2(model, identity_list, img_data, is_cuda=is_cuda)
         print('Accuracy: %f' % acc)

vectorizer/requirements.txt

@@ -1,4 +1,4 @@
 Flask
 Pillow
-https://download.pytorch.org/whl/cu100/torch-1.1.0-cp37-cp37m-linux_x86_64.whl
+torch
-https://download.pytorch.org/whl/cu100/torchvision-0.3.0-cp37-cp37m-linux_x86_64.whl
+torchvision

vectorizer/train-rec.sh

@@ -1,2 +1,2 @@
 python3 -m recognition.train --casia_list ~/datasets/CASIA-maxpy-clean/train.txt --casia_root ~/datasets/CASIA-maxpy-clean --lfw_root ~/datasets/lfw \
---lfw_pair_list lfw_test_pair.txt --model_name recongition1 --batch_size 20 --loss adacos --print_freq 20 --depth 50
+--lfw_pair_list lfw_test_pair.txt --model_name recongition1 --batch_size 20 --loss adacos --print_freq 20 --net resnet50

vectorizer/vectorizer/server.py

@@ -24,7 +24,7 @@ from flask import Flask, request, abort, jsonify
 from werkzeug.utils import secure_filename
 
 import torch
-from recognition.nets import get_net_by_depth
+from recognition.nets import get_net_by_name
 from torchvision import transforms as T
 from PIL import Image
 import identification.detector as fan
@@ -38,7 +38,7 @@ if fan_file is None:
 fan_model = fan.load_model(fan_file, is_cuda=is_cuda)
 
 # load recognition model
-rec_model = get_net_by_depth(int(os.environ.get('VS_REC_DEPTH', 50)))
+rec_model = get_net_by_name(os.environ.get('VS_REC_NET', 'resnet50'))
 rec_file = os.environ.get('VS_REC_MODEL', None)
 if rec_file is None:
     raise Exception('VS_REC_MODEL is mandatory parameter')