From 25ac9252879dc327d637ca409147193e27492559 Mon Sep 17 00:00:00 2001 From: Petr Masopust Date: Sat, 7 Sep 2019 21:25:54 +0200 Subject: [PATCH] ResNext models --- vectorizer/README.md | 6 +++--- vectorizer/recognition/nets.py | 38 ++++++++++++++++++++++++++------- vectorizer/recognition/test.py | 8 +++---- vectorizer/recognition/train.py | 11 +++++----- vectorizer/requirements.txt | 4 ++-- vectorizer/train-rec.sh | 2 +- vectorizer/vectorizer/server.py | 4 ++-- 7 files changed, 47 insertions(+), 26 deletions(-) diff --git a/vectorizer/README.md b/vectorizer/README.md index dd89b9a..0c55fa4 100644 --- a/vectorizer/README.md +++ b/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 | diff --git a/vectorizer/recognition/nets.py b/vectorizer/recognition/nets.py index d48fdf0..fb7f1aa 100644 --- a/vectorizer/recognition/nets.py +++ b/vectorizer/recognition/nets.py @@ -74,21 +74,43 @@ def sphere20(): return sphere20a() -def get_net_by_depth(depth): - if depth == 18: +def resnext50(pretrained=False, **kwargs): + """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 diff --git a/vectorizer/recognition/test.py b/vectorizer/recognition/test.py index b80b36b..d07063b 100644 --- a/vectorizer/recognition/test.py +++ b/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, - default=50) + parser.add_argument('--net', help='Net name, must be one of resnet18, resnet34, resnet50 (default), resnet101, resnet152, resnext50, resnext101 or spherenet', + 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) diff --git a/vectorizer/recognition/train.py b/vectorizer/recognition/train.py index 45dbf69..d8e9ea6 100644 --- a/vectorizer/recognition/train.py +++ b/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, - default=50) + parser.add_argument('--net', help='Net name, must be one of resnet18, resnet34, resnet50, resnet101, resnet152, resnext50, resnext101 or spherenet', + 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) diff --git a/vectorizer/requirements.txt b/vectorizer/requirements.txt index 71cd205..8cf5f7e 100644 --- a/vectorizer/requirements.txt +++ b/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 -https://download.pytorch.org/whl/cu100/torchvision-0.3.0-cp37-cp37m-linux_x86_64.whl +torch +torchvision diff --git a/vectorizer/train-rec.sh b/vectorizer/train-rec.sh index 53abbd9..41f3e13 100755 --- a/vectorizer/train-rec.sh +++ b/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 diff --git a/vectorizer/vectorizer/server.py b/vectorizer/vectorizer/server.py index 1589596..ff8033f 100644 --- a/vectorizer/vectorizer/server.py +++ b/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')