Change model to make it prunable

Author: daniil-lyakhov

torchreid/engine/engine.py

@@ -326,7 +326,18 @@ def run(
         self.fixbase_epoch = fixbase_epoch
         test_acc = AverageMeter()
         print('=> Start training')
-
+        res = self.test(
+                    self.epoch,
+                    dist_metric=dist_metric,
+                    normalize_feature=normalize_feature,
+                    visrank=visrank,
+                    visrank_topk=visrank_topk,
+                    save_dir=save_dir,
+                    use_metric_cuhk03=use_metric_cuhk03,
+                    ranks=ranks,
+                    lr_finder=lr_finder,
+                )
+        print(res)
         if perf_monitor and not lr_finder: perf_monitor.on_train_begin()
         for self.epoch in range(self.start_epoch, self.max_epoch):
             # change the NumPy’s seed at every epoch

torchreid/models/mobilenetv3.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
 
 from torchreid.losses import AngleSimpleLinear
 from torchreid.ops import Dropout, EvalModeSetter, rsc
@@ -33,19 +34,17 @@
 class SELayer(nn.Module):
     def __init__(self, channel, reduction=4):
         super(SELayer, self).__init__()
-        self.avg_pool = nn.AdaptiveAvgPool2d(1)
         self.fc = nn.Sequential(
-                nn.Linear(channel, make_divisible(channel // reduction, 8)),
+                nn.Conv2d(channel, make_divisible(channel // reduction, 8), 1),
                 nn.ReLU(inplace=True),
-                nn.Linear(make_divisible(channel // reduction, 8), channel),
+                nn.Conv2d(make_divisible(channel // reduction, 8), channel, 1),
                 HSigmoid()
         )
 
     def forward(self, x):
         with no_nncf_se_layer_context():
-            b, c, _, _ = x.size()
-            y = self.avg_pool(x).view(b, c)
-            y = self.fc(y).view(b, c, 1, 1)
+            y = F.adaptive_avg_pool2d(x, 1)
+            y = self.fc(y)
         return x * y
 
 

torchreid/utils/torchtools.py

@@ -285,6 +285,19 @@ def _print_loading_weights_inconsistencies(discarded_layers, unmatched_layers):
         )
 
 
+def update_checkpoint_mobilenet_v3(checkpoint):
+    fc = []
+    for k in checkpoint:
+        if 'fc' in k and not 'bias' in k:
+            fc.append(k)
+    for name in fc:
+        w = checkpoint[name]
+        shape = w.shape
+        w_new = w.view(shape + (1, 1))
+        print(name, ': ', checkpoint[name].shape, '->', w_new.shape)
+        checkpoint[name] = w_new
+
+
 def load_pretrained_weights(model, file_path='', pretrained_dict=None):
     r"""Loads pretrianed weights to model.
     Features::
@@ -317,6 +330,7 @@ def _remove_prefix(key, prefix):
     else:
         state_dict = checkpoint
 
+    update_checkpoint_mobilenet_v3(state_dict)
     model_dict = model.state_dict()
     new_state_dict = OrderedDict()
     matched_layers, discarded_layers = [], []