Change model to make it prunable

Author: daniil-lyakhov

torchreid/apis/training.py

@@ -31,12 +31,10 @@
 
 def run_lr_finder(cfg, datamanager, model, optimizer, scheduler, classes,
                   rebuild_model=True, gpu_num=1, split_models=False):
-    if rebuild_model:
-        tmp_model = model
-    else:
-        tmp_model = deepcopy(model)
+    if not rebuild_model:
+        backup_model = deepcopy(model)
 
-    engine = build_engine(cfg, datamanager, tmp_model, optimizer, scheduler, initial_lr=cfg.train.lr)
+    engine = build_engine(cfg, datamanager, model, optimizer, scheduler, initial_lr=cfg.train.lr)
     lr_finder = LrFinder(engine=engine, **lr_finder_run_kwargs(cfg))
     aux_lr = lr_finder.process()
 
@@ -54,16 +52,18 @@ def run_lr_finder(cfg, datamanager, model, optimizer, scheduler, classes,
     set_random_seed(cfg.train.seed, cfg.train.deterministic)
     datamanager = build_datamanager(cfg, classes)
     num_train_classes = datamanager.num_train_pids
+
     if rebuild_model:
-        model = torchreid.models.build_model(**model_kwargs(cfg, num_train_classes))
+        backup_model = torchreid.models.build_model(**model_kwargs(cfg, num_train_classes))
         num_aux_models = len(cfg.mutual_learning.aux_configs)
-        model, _ = put_main_model_on_the_device(model, cfg.use_gpu, gpu_num, num_aux_models, split_models)
-    optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
+        backup_model, _ = put_main_model_on_the_device(backup_model, cfg.use_gpu, gpu_num, num_aux_models, split_models)
+
+    optimizer = torchreid.optim.build_optimizer(backup_model, **optimizer_kwargs(cfg))
     scheduler = torchreid.optim.build_lr_scheduler(optimizer=optimizer,
                                                    num_iter=datamanager.num_iter,
                                                    **lr_scheduler_kwargs(cfg))
 
-    return cfg.train.lr
+    return cfg.train.lr, backup_model, optimizer, scheduler
 
 
 def run_training(cfg, datamanager, model, optimizer, scheduler, extra_device_ids, init_lr,

torchreid/engine/engine.py

@@ -326,7 +326,6 @@ def run(
         self.fixbase_epoch = fixbase_epoch
         test_acc = AverageMeter()
         print('=> Start training')
-
         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 = [], []

training/pruning_int8/mobilenet_v3_small/compress.sh

@@ -0,0 +1,9 @@
+python /home/automation/dlyakhov/training_extensions/external/deep-object-reid/tools/main.py \
+--config-file /home/automation/dlyakhov/training_extensions/external/deep-object-reid/training/pruning_int8/mobilenet_v3_small/mobilenetv3_small.yml \
+--gpu-num 1 \
+--custom-roots \
+/home/automation/dlyakhov/training_extensions/external/deep-object-reid/CIFAR100/train \
+/home/automation/dlyakhov/training_extensions/external/deep-object-reid/CIFAR100/val \
+--root _ \
+model.load_weights /mnt/icv_externalN/dlyakhov/ote-classification-checkpoint/CIFAR100_mobielenet_v3_small/model_0/{name}.pth.tar-142 \
+data.save_dir /home/automation/dlyakhov/training_extensions/external/deep-object-reid/training/pruning_int8/mobilenet_v3_small/output_21_09

training/pruning_int8/mobilenet_v3_small/compress_by_flops.sh

@@ -0,0 +1,9 @@
+python /home/automation/dlyakhov/training_extensions/external/deep-object-reid/tools/main.py \
+--config-file /home/automation/dlyakhov/training_extensions/external/deep-object-reid/training/pruning_int8/mobilenet_v3_small/mobilenetv3_small_by_flops.yml \
+--gpu-num 1 \
+--custom-roots \
+/home/automation/dlyakhov/training_extensions/external/deep-object-reid/CIFAR100/train \
+/home/automation/dlyakhov/training_extensions/external/deep-object-reid/CIFAR100/val \
+--root _ \
+model.load_weights /mnt/icv_externalN/dlyakhov/ote-classification-checkpoint/CIFAR100_mobielenet_v3_small/model_0/{name}.pth.tar-142 \
+data.save_dir /home/automation/dlyakhov/training_extensions/external/deep-object-reid/training/pruning_int8/mobilenet_v3_small/output_pruning_0.1_filters_by_flops_13_08

training/pruning_int8/mobilenet_v3_small/compress_sparse.sh

@@ -6,4 +6,4 @@ python /home/automation/dlyakhov/training_extensions/external/deep-object-reid/t
 /home/automation/dlyakhov/training_extensions/external/deep-object-reid/CIFAR100/val \
 --root _ \
 model.load_weights /mnt/icv_externalN/dlyakhov/ote-classification-checkpoint/CIFAR100_mobielenet_v3_small/model_0/{name}.pth.tar-142 \
-data.save_dir /home/automation/dlyakhov/training_extensions/external/deep-object-reid/training/pruning_int8/mobilenet_v3_small/sparsity_magnitude_24_09
+data.save_dir /home/automation/dlyakhov/training_extensions/external/deep-object-reid/training/pruning_int8/mobilenet_v3_small/sparsity_23_09

training/pruning_int8/mobilenet_v3_small/mobilenetv3_large_aux.yml

@@ -35,8 +35,10 @@ train:
   deterministic: True
   patience: 5
   gamma: 0.1
-  sam:
-    rho: 0.05
+
+sam:
+  enable: True
+  rho: 0.05
 
 test:
   batch_size: 128

training/pruning_int8/mobilenet_v3_small/mobilenetv3_small.yml

@@ -0,0 +1,92 @@
+
+lr_finder:
+  enable: True
+  mode: TPE
+  stop_after: False
+  num_epochs: 6
+  step: 0.001
+  epochs_warmup: 1
+  path_to_savefig: 'lr_finder.jpg'
+  max_lr: 0.029
+  min_lr: 0.005
+  n_trials: 15
+
+model:
+  name: 'mobilenetv3_small'
+  type: 'classification'
+  pretrained: True
+  save_chkpt: True
+  feature_dim: 1024
+
+mutual_learning:
+  aux_configs: ['mobilenetv3_large_aux.yml']
+
+custom_datasets:
+  roots: ['data/CIFAR100/train', 'data/CIFAR100/val']
+  types: ['classification_image_folder', 'classification_image_folder']
+  names: ['CIFAR100_train', 'CIFAR100_val']
+
+data:
+  root: './'
+  sources: ['CIFAR100_train']
+  targets: ['CIFAR100_val']
+  height: 224
+  width: 224
+  norm_mean: [0.485, 0.456, 0.406]
+  norm_std: [0.229, 0.224, 0.225]
+  save_dir: 'output/mobilenetv3_small/log'
+  workers: 6
+  transforms:
+    random_flip:
+      enable: True
+      p: 0.5
+    augmix:
+      enable: True
+      cfg_str: "augmix-m5-w3"
+
+loss:
+  name: 'softmax'
+  softmax:
+    s: 1.0
+    compute_s: False
+
+sampler:
+  train_sampler: 'RandomSampler'
+
+metric_losses:
+  enable: False
+
+train:
+  optim: 'sam'
+  lr: 0.013
+  nbd: True
+  max_epoch: 200
+  weight_decay: 5e-4
+  batch_size: 84
+  lr_scheduler: 'warmup'
+  warmup: 5
+  base_scheduler: 'reduce_on_plateau_delayed'
+  epoch_delay: 40
+  early_stoping: True
+  train_patience: 5
+  lr_decay_factor: 200
+  deterministic: True
+  patience: 5
+  gamma: 0.1
+  ema:
+    enable: True
+    ema_decay: 0.999
+
+sam:
+  enable: True
+  rho: 0.05
+
+test:
+  batch_size: 128
+  evaluate: False
+  eval_freq: 1
+
+nncf:
+  enable: True
+  coeff_decrease_lr_for_nncf: 1.
+  nncf_config_path: 'nncf_config.json'

training/pruning_int8/mobilenet_v3_small/mobilenetv3_small_by_flops.yml

@@ -0,0 +1,91 @@
+
+lr_finder:
+  enable: True
+  mode: TPE
+  stop_after: False
+  num_epochs: 6
+  step: 0.001
+  epochs_warmup: 1
+  path_to_savefig: 'lr_finder.jpg'
+  max_lr: 0.029
+  min_lr: 0.005
+  n_trials: 15
+
+model:
+  name: 'mobilenetv3_small'
+  type: 'classification'
+  pretrained: True
+  save_chkpt: True
+  feature_dim: 1024
+
+mutual_learning:
+  aux_configs: ['mobilenetv3_large_aux.yml']
+
+custom_datasets:
+  roots: ['data/CIFAR100/train', 'data/CIFAR100/val']
+  types: ['classification_image_folder', 'classification_image_folder']
+  names: ['CIFAR100_train', 'CIFAR100_val']
+
+data:
+  root: './'
+  sources: ['CIFAR100_train']
+  targets: ['CIFAR100_val']
+  height: 224
+  width: 224
+  norm_mean: [0.485, 0.456, 0.406]
+  norm_std: [0.229, 0.224, 0.225]
+  save_dir: 'output/mobilenetv3_small/log'
+  workers: 6
+  transforms:
+    random_flip:
+      enable: True
+      p: 0.5
+    augmix:
+      enable: True
+      cfg_str: "augmix-m5-w3"
+
+loss:
+  name: 'softmax'
+  softmax:
+    s: 1.0
+    compute_s: False
+
+sampler:
+  train_sampler: 'RandomSampler'
+
+metric_losses:
+  enable: False
+
+train:
+  optim: 'sam'
+  lr: 0.013
+  nbd: True
+  max_epoch: 200
+  weight_decay: 5e-4
+  batch_size: 84
+  lr_scheduler: 'warmup'
+  warmup: 5
+  base_scheduler: 'reduce_on_plateau'
+  early_stoping: True
+  train_patience: 5
+  lr_decay_factor: 200
+  deterministic: True
+  patience: 5
+  gamma: 0.1
+  ema:
+    enable: True
+    ema_decay: 0.999
+
+sam:
+  enable: True
+  rho: 0.05
+
+test:
+  batch_size: 128
+  evaluate: False
+  eval_freq: 1
+
+nncf:
+  enable: True
+  coeff_decrease_lr_for_nncf: 1.
+  nncf_config_path: 'nncf_config_by_flops.json'