Create MetaBNNeck

ppcls/arch/gears/__init__.py

@@ -21,6 +21,7 @@ from paddle.nn import Tanh
 from .bnneck import BNNeck
 from .adamargin import AdaMargin
 from .frfn_neck import FRFNNeck
+from .metabnneck import MetaBNNeck
 
 __all__ = ['build_gear']
 
@@ -28,7 +29,7 @@ __all__ = ['build_gear']
 def build_gear(config):
     support_dict = [
         'ArcMargin', 'CosMargin', 'CircleMargin', 'FC', 'VehicleNeck', 'Tanh',
-        'BNNeck', 'AdaMargin', 'FRFNNeck'
+        'BNNeck', 'AdaMargin', 'FRFNNeck', 'MetaBNNeck'
     ]
     module_name = config.pop('name')
     assert module_name in support_dict, Exception(

ppcls/arch/gears/metabnneck.py

+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import, division, print_function
+
+from collections import defaultdict
+import copy
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from ..utils import get_param_attr_dict
+
+
+class MetaBN1D(nn.BatchNorm2D):
+    def forward(self, inputs, opt={}):
+        mode = opt.get("bn_mode", "general") if self.training else "eval"
+        if mode == "general":  # update, but not apply running_mean/var
+            result = F.batch_norm(inputs, self._mean, self._variance,
+                                  self.weight, self.bias, self.training,
+                                  self._momentum, self._epsilon)
+        elif mode == "hold":  # not update, not apply running_mean/var
+            result = F.batch_norm(
+                inputs,
+                paddle.mean(
+                    inputs, axis=0),
+                paddle.var(inputs, axis=0),
+                self.weight,
+                self.bias,
+                self.training,
+                self._momentum,
+                self._epsilon)
+        elif mode == "eval":  # fix and apply running_mean/var,
+            if self._mean is None:
+                result = F.batch_norm(
+                    inputs,
+                    paddle.mean(
+                        inputs, axis=0),
+                    paddle.var(inputs, axis=0),
+                    self.weight,
+                    self.bias,
+                    True,
+                    self._momentum,
+                    self._epsilon)
+            else:
+                result = F.batch_norm(inputs, self._mean, self._variance,
+                                      self.weight, self.bias, False,
+                                      self._momentum, self._epsilon)
+        return result
+
+
+class MetaBNNeck(nn.Layer):
+    def __init__(self, num_features, **kwargs):
+        super(MetaBNNeck, self).__init__()
+        weight_attr = paddle.ParamAttr(
+            initializer=paddle.nn.initializer.Constant(value=1.0))
+        bias_attr = paddle.ParamAttr(
+            initializer=paddle.nn.initializer.Constant(value=0.0),
+            trainable=False)
+
+        if 'weight_attr' in kwargs:
+            weight_attr = get_param_attr_dict(kwargs['weight_attr'])
+
+        bias_attr = None
+        if 'bias_attr' in kwargs:
+            bias_attr = get_param_attr_dict(kwargs['bias_attr'])
+
+        use_global_stats = None
+        if 'use_global_stats' in kwargs:
+            use_global_stats = get_param_attr_dict(kwargs['use_global_stats'])
+
+        self.feat_bn = MetaBN1D(
+            num_features,
+            momentum=0.9,
+            epsilon=1e-05,
+            weight_attr=weight_attr,
+            bias_attr=bias_attr,
+            use_global_stats=use_global_stats)
+        self.flatten = nn.Flatten()
+        self.opt = {}
+
+    def forward(self, x):
+        x = self.flatten(x)
+        x = self.feat_bn(x, self.opt)
+        return x
+
+    def reset_opt(self):
+        self.opt = defaultdict()
+
+    def setup_opt(self, opt):
+        """
+        enable_inside_update: enable inside updating for `gate` in MetaBIN
+        lr_gate: learning rate of `gate` during meta-train phase
+        bn_mode: control the running stats & updating of BN
+        """
+        self.check_opt(opt)
+        self.opt = copy.deepcopy(opt)
+
+    @classmethod
+    def check_opt(cls, opt):
+        assert isinstance(opt, dict), \
+            TypeError('Got the wrong type of `opt`. Please use `dict` type.')
+
+        if opt.get('enable_inside_update', False) and 'lr_gate' not in opt:
+            raise RuntimeError('Missing `lr_gate` in opt.')
+
+        assert isinstance(opt.get('lr_gate', 1.0), float), \
+            TypeError('Got the wrong type of `lr_gate`. Please use `float` type.')
+        assert isinstance(opt.get('enable_inside_update', True), bool), \
+            TypeError('Got the wrong type of `enable_inside_update`. Please use `bool` type.')
+        assert opt.get('bn_mode', "general") in ["general", "hold", "eval"], \
+            TypeError('Got the wrong value of `bn_mode`.')

ppcls/configs/reid/MetaBIN_ResNet50_single_source.yaml

@@ -34,7 +34,7 @@ Arch:
   BackboneStopLayer:
     name: "flatten"
   Neck:
-    name: BNNeck
+    name: MetaBNNeck
     num_features: &feat_dim 2048
     use_global_stats: True
   Head:

ppcls/engine/train/train_metabin.py

@@ -99,7 +99,7 @@ def train_epoch_metabin(engine, epoch_id, print_batch_step):
         }
         # step lr (by iter)
         # the last lr_sch is cyclic_lr
-        for i in range(len(engine.lr_sch) - 1):
+        for i in range(len(engine.lr_sch)):
             if not getattr(engine.lr_sch[i], "by_epoch", False):
                 engine.lr_sch[i].step()
         # update ema
@@ -118,7 +118,7 @@ def train_epoch_metabin(engine, epoch_id, print_batch_step):
 
     # step lr(by epoch)
     # the last lr_sch is cyclic_lr
-    for i in range(len(engine.lr_sch) - 1):
+    for i in range(len(engine.lr_sch)):
         if getattr(engine.lr_sch[i], "by_epoch", False) and \
                 type_name(engine.lr_sch[i]) != "ReduceOnPlateau":
             engine.lr_sch[i].step()
@@ -138,19 +138,20 @@ def setup_opt(engine, stage):
     elif stage == "mtest":
         norm_lr = engine.lr_sch[1].last_lr
         cyclic_lr = engine.lr_sch[2].get_lr()
-        engine.lr_sch[2].step()  # update cyclic learning rate
         opt["bn_mode"] = "hold"
         opt["enable_inside_update"] = True
         opt["lr_gate"] = norm_lr * cyclic_lr
     for name, layer in engine.model.backbone.named_sublayers():
         if "bn" == name.split('.')[-1]:
             layer.setup_opt(opt)
+    engine.model.neck.setup_opt(opt)
 
 
 def reset_opt(model):
     for name, layer in model.backbone.named_sublayers():
         if "bn" == name.split('.')[-1]:
             layer.reset_opt()
+    model.neck.reset_opt()
 
 
 def get_meta_data(meta_dataloader_iter, num_domain):
@@ -200,6 +201,7 @@ def forward(engine, batch, loss_func):
 
 
 def backward(engine, loss, optimizer):
+    optimizer.clear_grad()
     scaled = engine.scaler.scale(loss)
     scaled.backward()
     engine.scaler.minimize(optimizer, scaled)