Update face recognition scripts (openvinotoolkit#134)

* Face recognition: make all models accessible via config

* Update face recognition readme

* Face recognition: add new congig, fixes in scripts

* Add link to mobilenetv2_2x

Author: sovrasov

pytorch_toolkit/face_recognition/README.md

@@ -36,14 +36,16 @@ cd $FR_ROOT/
 ```
 
 2. To start training FR model:
+
 ```bash
 python train.py --train_data_root $VGGFace2_ROOT/train/ --train_list $VGGFace2_ROOT/meta/train_list.txt
 --train_landmarks  $VGGFace2_ROOT/bb_landmark/ --val_data_root  $LFW_ROOT/lfw/ --val_list $LFW_ROOT/pairs.txt  
 --val_landmarks $LFW_ROOT/lfw_landmark.txt --train_batch_size 200  --snap_prefix mobilenet_256 --lr 0.35
---embed_size 256 --model mobilenet --device 1
+--embed_size 256 --model mobilenetv2 --device 1
 ```
 
 3. To evaluate FR snapshot (let's say we have MobileNet with 256 embedding size trained for 300k):
+
 ```bash
  python evaluate_lfw.py --val_data_root $LFW_ROOT/lfw/ --val_list $LFW_ROOT/pairs.txt
  --val_landmarks $LFW_ROOT/lfw_landmark.txt --snap /path/to/snapshot/mobilenet_256_300000.pt --model mobilenet --embed_size 256
@@ -62,7 +64,7 @@ margin_type: cos
 s: 30
 m: 0.35
 #model parameters
-model: mobilenet
+model: mobilenetv2
 embed_size: 256
 #misc
 snap_prefix: MobileFaceNet
@@ -81,14 +83,17 @@ python train.py -m 0.35 @./my_config.yml #here m can be overwritten with the val
 
 ## Models
 
-1. You can download pretrained model from fileshare as well - https://download.01.org/opencv/openvino_training_extensions/models/face_recognition/Mobilenet_se_focal_121000.pt
+1. You can download pretrained model from fileshare as well - [mobilenetv2](https://download.01.org/opencv/openvino_training_extensions/models/face_recognition/Mobilenet_se_focal_121000.pt),
+[mobilenetv2_2x](https://download.01.org/opencv/openvino_training_extensions/models/face_recognition/Mobilenet_2x_se_121000.pt).
+
 ```bash
 cd $FR_ROOT
 python evaluate_lfw.py --val_data_root $LFW_ROOT/lfw/ --val_list $LFW_ROOT/pairs.txt --val_landmarks $LFW_ROOT/lfw_landmark.txt
 --snap /path/to/snapshot/Mobilenet_se_focal_121000.pt --model mobilenet --embed_size 256
 ```
 
 2. You should get the following output:
+- for `mobilenetv2`:
 ```
 I1114 09:33:37.846870 10544 evaluate_lfw.py:242] Accuracy/Val_same_accuracy mean: 0.9923
 I1114 09:33:37.847019 10544 evaluate_lfw.py:243] Accuracy/Val_diff_accuracy mean: 0.9970
@@ -97,6 +102,18 @@ I1114 09:33:37.847179 10544 evaluate_lfw.py:245] Accuracy/Val_accuracy std dev:
 I1114 09:33:37.847229 10544 evaluate_lfw.py:246] AUC: 0.9995
 I1114 09:33:37.847305 10544 evaluate_lfw.py:247] Estimated threshold: 0.7241
 ```
+- for `mobilenetv2_2x`:
+```
+I0820 15:48:06.307454 23328 evaluate_lfw.py:262] Accuracy/Val_same_accuracy mean: 0.9893
+I0820 15:48:06.307612 23328 evaluate_lfw.py:263] Accuracy/Val_diff_accuracy mean: 0.9990
+I0820 15:48:06.307647 23328 evaluate_lfw.py:264] Accuracy/Val_accuracy mean: 0.9942
+I0820 15:48:06.307732 23328 evaluate_lfw.py:265] Accuracy/Val_accuracy std dev: 0.0061
+I0820 15:48:06.307766 23328 evaluate_lfw.py:266] AUC: 0.9992
+I0820 15:48:06.307812 23328 evaluate_lfw.py:267] Estimated threshold: 0.6721
+```
+
+`mobilenetv2_2x` is slightly worse on the LFW benchmark than `mobilenetv2`, but it's heavier and achieves higher score in the
+uncleaned version of the [MegaFace](http://megaface.cs.washington.edu/participate/challenge.html) benchmark: 73.77% rank-1 at 1M distractors in reidentification protocol vs 70.2%.
 
 ## Face Recognition Demo
 

pytorch_toolkit/face_recognition/configs/mobilefacenet_2x_vgg2.yml

@@ -0,0 +1,16 @@
+#optimizer parameters
+lr: 0.4
+train_batch_size: 256
+#loss options
+margin_type: cos
+s: 30
+m: 0.35
+mining_type: sv
+t: 1.1
+#model parameters
+model: mobilenetv2_2x
+embed_size: 256
+
+train_dataset: vgg
+snap_prefix: MobileFaceNet
+devices: [0, 1]

pytorch_toolkit/face_recognition/configs/mobilefacenet_vgg2.yml

@@ -8,7 +8,7 @@ m: 0.35
 mining_type: sv
 t: 1.1
 #model parameters
-model: mobilenet
+model: mobilenetv2
 embed_size: 256
 
 train_dataset: vgg

pytorch_toolkit/face_recognition/dump_features.py

@@ -120,7 +120,7 @@ def main(args):
 
     emb_array = np.zeros((nrof_images, args.embedding_size), dtype=np.float32)
 
-    dataset.transform = t.Compose([ResizeNumpy(models_backbones[args.model].get_input_res()),
+    dataset.transform = t.Compose([ResizeNumpy(models_backbones[args.model]().get_input_res()),
                                    NumpyToTensor(switch_rb=True)])
     val_loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=5, shuffle=False)
 

pytorch_toolkit/face_recognition/evaluate_lfw.py

@@ -271,7 +271,7 @@ def evaluate(args, dataset, model, compute_embeddings_fun, val_batch_size=16,
 
 def load_test_dataset(arguments):
     """Loads and configures the LFW dataset"""
-    input_size = models_backbones[arguments.model].get_input_res()
+    input_size = models_backbones[arguments.model]().get_input_res()
     lfw = LFW(arguments.val, arguments.v_list, arguments.v_land)
     assert lfw.use_landmarks
     log.info('Using landmarks for the LFW images.')

pytorch_toolkit/face_recognition/model/backbones/resnet.py

@@ -18,7 +18,7 @@
 
 
 class ResNet(nn.Module):
-    def __init__(self, block, layers, num_classes=1000, activation=nn.ReLU):
+    def __init__(self, block, layers, num_classes=1000, activation=nn.ReLU, head=False):
         self.inplanes = 64
         super(ResNet, self).__init__()
         self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1,
@@ -32,7 +32,12 @@ def __init__(self, block, layers, num_classes=1000, activation=nn.ReLU):
         self.layer4 = self._make_layer(block, 512, layers[3], stride=2, activation=activation)
         self.avgpool = nn.Conv2d(512 * block.expansion, 512 * block.expansion, 7,
                                  groups=512 * block.expansion, bias=False)
-        self.fc = nn.Conv2d(512 * block.expansion, num_classes, 1, stride=1, padding=0, bias=False)
+        self.head = head
+        if not self.head:
+            self.output_channels = 512 * block.expansion
+        else:
+            self.fc = nn.Conv2d(512 * block.expansion, num_classes, 1, stride=1, padding=0, bias=False)
+            self.output_channels = num_classes
 
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
@@ -70,10 +75,14 @@ def forward(self, x):
         x = self.layer4(x)
 
         x = self.avgpool(x)
-        x = self.fc(x)
+        if self.head:
+            x = self.fc(x)
 
         return x
 
+    def get_output_channels(self):
+        return self.output_channels
+
 
 def resnet50(**kwargs):
     model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)

pytorch_toolkit/face_recognition/model/backbones/se_resnet.py

@@ -12,16 +12,15 @@
 """
 
 import math
-
 import torch.nn as nn
 
 from model.blocks.se_resnet_blocks import SEBottleneck
 
 
 class SEResNet(nn.Module):
-    def __init__(self, block, layers, num_classes=1000, activation=nn.ReLU):
-        self.inplanes = 64
+    def __init__(self, block, layers, num_classes=1000, activation=nn.ReLU, head=False):
         super(SEResNet, self).__init__()
+        self.inplanes = 64
         self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1,
                                bias=False)
         self.bn1 = nn.BatchNorm2d(64)
@@ -33,7 +32,12 @@ def __init__(self, block, layers, num_classes=1000, activation=nn.ReLU):
         self.layer4 = self._make_layer(block, 512, layers[3], stride=2, activation=activation)
         self.avgpool = nn.Conv2d(512 * block.expansion, 512 * block.expansion, 7,
                                  groups=512 * block.expansion, bias=False)
-        self.fc = nn.Conv2d(512 * block.expansion, num_classes, 1, stride=1, padding=0, bias=False)
+        self.head = head
+        if not self.head:
+            self.output_channels = 512 * block.expansion
+        else:
+            self.fc = nn.Conv2d(512 * block.expansion, num_classes, 1, stride=1, padding=0, bias=False)
+            self.output_channels = num_classes
 
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
@@ -72,10 +76,14 @@ def forward(self, x):
         x = self.layer4(x)
 
         x = self.avgpool(x)
-        x = self.fc(x)
+        if self.head:
+            x = self.fc(x)
 
         return x
 
+    def get_output_channels(self):
+        return self.output_channels
+
 
 def se_resnet50(**kwargs):
     model = SEResNet(SEBottleneck, [3, 4, 6, 3], **kwargs)

pytorch_toolkit/face_recognition/model/backbones/se_resnext.py

@@ -18,25 +18,29 @@
 
 
 class SEResNeXt(nn.Module):
-
-    def __init__(self, block, layers, cardinality=32, num_classes=1000):
+    def __init__(self, block, layers, cardinality=32, num_classes=1000, activation=nn.ReLU, head=False):
         super(SEResNeXt, self).__init__()
         self.cardinality = cardinality
         self.inplanes = 64
 
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1, padding=3,
                                bias=False)
         self.bn1 = nn.BatchNorm2d(64)
         self.relu = nn.ReLU(inplace=True)
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
 
-        self.layer1 = self._make_layer(block, 64, layers[0])
-        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
-        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
-
-        self.avgpool = nn.AdaptiveAvgPool2d(1)
-        self.fc = nn.Linear(512 * block.expansion, num_classes)
+        self.layer1 = self._make_layer(block, 64, layers[0], activation=activation)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, activation=activation)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, activation=activation)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, activation=activation)
+        self.avgpool = nn.Conv2d(512 * block.expansion, 512 * block.expansion, 7,
+                                 groups=512 * block.expansion, bias=False)
+        self.head = head
+        if not self.head:
+            self.output_channels = 512 * block.expansion
+        else:
+            self.fc = nn.Conv2d(512 * block.expansion, num_classes, 1, stride=1, padding=0, bias=False)
+            self.output_channels = num_classes
 
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
@@ -48,7 +52,7 @@ def __init__(self, block, layers, cardinality=32, num_classes=1000):
                 m.weight.data.fill_(1)
                 m.bias.data.zero_()
 
-    def _make_layer(self, block, planes, blocks, stride=1):
+    def _make_layer(self, block, planes, blocks, stride=1, activation=nn.ReLU):
         downsample = None
         if stride != 1 or self.inplanes != planes * block.expansion:
             downsample = nn.Sequential(
@@ -58,10 +62,10 @@ def _make_layer(self, block, planes, blocks, stride=1):
             )
 
         layers = []
-        layers.append(block(self.inplanes, planes, self.cardinality, stride, downsample))
+        layers.append(block(self.inplanes, planes, self.cardinality, stride, downsample, activation=activation))
         self.inplanes = planes * block.expansion
         for _ in range(1, blocks):
-            layers.append(block(self.inplanes, planes, self.cardinality))
+            layers.append(block(self.inplanes, planes, self.cardinality, activation=activation))
 
         return nn.Sequential(*layers)
 
@@ -77,12 +81,14 @@ def forward(self, x):
         x = self.layer4(x)
 
         x = self.avgpool(x)
-        x = x.view(x.size(0), -1)
-
-        x = self.fc(x)
+        if self.head:
+            x = self.fc(x)
 
         return x
 
+    def get_output_channels(self):
+        return self.output_channels
+
 
 def se_resnext50(**kwargs):
     model = SEResNeXt(SEBottleneckX, [3, 4, 6, 3], **kwargs)

pytorch_toolkit/face_recognition/model/blocks/se_resnet_blocks.py

@@ -30,7 +30,10 @@ def __init__(self, inplanes, planes, stride=1, downsample=None, activation=nn.Re
         self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
         self.bn3 = nn.BatchNorm2d(planes * 4)
 
-        self.relu = make_activation(activation)
+        self.relu1 = make_activation(activation)
+        self.relu2 = make_activation(activation)
+        self.relu3 = make_activation(activation)
+        self.relu4 = make_activation(activation)
 
         # SE
         self.global_pool = nn.AdaptiveAvgPool2d(1)
@@ -47,25 +50,25 @@ def forward(self, x):
 
         out = self.conv1(x)
         out = self.bn1(out)
-        out = self.relu(out)
+        out = self.relu1(out)
 
         out = self.conv2(out)
         out = self.bn2(out)
-        out = self.relu(out)
+        out = self.relu2(out)
 
         out = self.conv3(out)
         out = self.bn3(out)
 
         out1 = self.global_pool(out)
         out1 = self.conv_down(out1)
-        out1 = self.relu(out1)
+        out1 = self.relu3(out1)
         out1 = self.conv_up(out1)
         out1 = self.sig(out1)
 
         if self.downsample is not None:
             residual = self.downsample(x)
 
         res = out1 * out + residual
-        res = self.relu(res)
+        res = self.relu4(res)
 
         return res

pytorch_toolkit/face_recognition/model/blocks/se_resnext_blocks.py

@@ -14,12 +14,13 @@
 import torch.nn as nn
 
 from model.blocks.shared_blocks import SELayer
+from model.blocks.shared_blocks import make_activation
 
 
 class SEBottleneckX(nn.Module):
     expansion = 4
 
-    def __init__(self, inplanes, planes, cardinality, stride=1, downsample=None):
+    def __init__(self, inplanes, planes, cardinality, stride=1, downsample=None, activation=nn.ReLU):
         super(SEBottleneckX, self).__init__()
         self.conv1 = nn.Conv2d(inplanes, planes * 2, kernel_size=1, bias=False)
         self.bn1 = nn.BatchNorm2d(planes * 2)
@@ -31,9 +32,12 @@ def __init__(self, inplanes, planes, cardinality, stride=1, downsample=None):
         self.conv3 = nn.Conv2d(planes * 2, planes * 4, kernel_size=1, bias=False)
         self.bn3 = nn.BatchNorm2d(planes * 4)
 
-        self.selayer = SELayer(planes * 4, 16, nn.ReLU)
+        self.selayer = SELayer(planes * 4, 16, activation)
+
+        self.relu1 = make_activation(activation)
+        self.relu2 = make_activation(activation)
+        self.relu3 = make_activation(activation)
 
-        self.relu = nn.ReLU(inplace=True)
         self.downsample = downsample
         self.stride = stride
 
@@ -42,11 +46,11 @@ def forward(self, x):
 
         out = self.conv1(x)
         out = self.bn1(out)
-        out = self.relu(out)
+        out = self.relu1(out)
 
         out = self.conv2(out)
         out = self.bn2(out)
-        out = self.relu(out)
+        out = self.relu2(out)
 
         out = self.conv3(out)
         out = self.bn3(out)
@@ -57,6 +61,6 @@ def forward(self, x):
             residual = self.downsample(x)
 
         out += residual
-        out = self.relu(out)
+        out = self.relu3(out)
 
         return out