Add LocalResponseNorm and BatchNorm in training

scripts/train_bottleneck.py

@@ -33,7 +33,8 @@
         {'head': 50, 'bottlenecks': 30, 'full': 50},
     'valid-split': 0.1,
     'recompute-bottlenecks': True,
-    'generator': {}
+    'generator': {},
+    'use_lrn': True,
 }
 NAME_BOTTLENECKS = 'bottlenecks.npz'
 NAME_CHECKPOINT = 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5'
@@ -51,7 +52,7 @@ def _main(path_dataset, path_anchors, path_weights=None, path_output='.',
     # make sure you know what you freeze
     model, bottleneck_model, last_layer_model = create_model_bottleneck(
         config['image-size'], anchors, nb_classes, freeze_body=2,
-        weights_path=path_weights, nb_gpu=nb_gpu)
+        weights_path=path_weights, nb_gpu=nb_gpu, lrn=config['use_lrn'])
 
     log_tb = TensorBoard(log_dir=path_output)
     checkpoint = ModelCheckpoint(os.path.join(path_output, NAME_CHECKPOINT),

scripts/training.py

@@ -35,6 +35,7 @@
     'epochs':
         {'head': 50, 'full': 50},
     'valid-split': 0.1,
+    'use_lrn': True,
     'generator': {
         'jitter': 0.3,
         'color_hue': 0.1,
@@ -75,6 +76,7 @@ def load_config(path_config, default_config):
         return config
     with open(path_config, 'r') as fp:
         conf_user = yaml.safe_load(fp)
+
     config.update(conf_user)
     return config
 
@@ -123,6 +125,7 @@ def _main(path_dataset, path_anchors, path_weights=None, path_output='.',
           path_config=None, path_classes=None, nb_gpu=1, **kwargs):
 
     config = load_config(path_config, DEFAULT_CONFIG)
+
     anchors = get_anchors(path_anchors)
 
     nb_classes = get_nb_classes(path_dataset)
@@ -133,7 +136,7 @@ def _main(path_dataset, path_anchors, path_weights=None, path_output='.',
     _create_model = create_model_tiny if is_tiny_version else create_model
     name_prefix = 'tiny-' if is_tiny_version else ''
     model = _create_model(config['image-size'], anchors, nb_classes, freeze_body=2,
-                          weights_path=path_weights, nb_gpu=nb_gpu)
+                          weights_path=path_weights, nb_gpu=nb_gpu, lrn=config['use_lrn'])
 
     tb_logging = TensorBoard(log_dir=path_output)
     checkpoint = ModelCheckpoint(os.path.join(path_output, NAME_CHECKPOINT),

yolo3/local_response_norm.py

@@ -0,0 +1,118 @@
+from keras.layers.core import Layer
+from keras import backend as K
+
+
+class LocalResponseNorm(Layer):
+    """
+    LocalResponseNorm Layer Definition in Keras
+
+    :param float alpha:
+    :param int k:
+    :param float beta:
+    :param int n:
+    :param **kwargs of Layer Class from Keras:
+
+    This code is adapted from pylearn2.
+    License at: https://github.com/lisa-lab/pylearn2/blob/master/LICENSE.txt
+    """
+    def __init__(self, alpha=0.0001, k=1, beta=0.75, n=5, **kwargs):
+        self.alpha = alpha
+        self.k = k
+        self.beta = beta
+        self.n = n
+        super(LocalResponseNorm, self).__init__(**kwargs)
+
+    def call(self, x, mask=None):
+        b, ch, r, c = x.shape
+        half_n = self.n // 2
+        input_sqr = K.square(x)
+
+        extra_channels = K.zeros((b, int(ch) + 2 * half_n, r, c))
+        input_sqr = K.concatenate(
+                                [
+                                    extra_channels[:, :half_n, :, :],
+                                    input_sqr,
+                                    extra_channels[:, half_n + int(ch):, :, :]
+                                ],
+                                axis=1
+                                )
+
+        scale = self.k  # offset for the scale
+        norm_alpha = self.alpha / self.n  # normalized alpha
+        for i in range(self.n):
+            scale += norm_alpha * input_sqr[:, i:i + int(ch), :, :]
+        scale = scale ** self.beta
+        x = x / scale
+        return x
+
+    def get_config(self):
+        config = {"alpha": self.alpha,
+                  "k": self.k,
+                  "beta": self.beta,
+                  "n": self.n}
+        base_config = super(LocalResponseNorm, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+class LocalResponseNorm2D(LocalResponseNorm):
+    """
+    LocalResponseNorm2D Layer Definition in Keras
+
+    :param floatalpha:
+    :param int k:
+    :param float beta:
+    :param int n:
+    :param **kwargs of Layer Class from Keras:
+
+    This code is adapted from pylearn2.
+    License at: https://github.com/lisa-lab/pylearn2/blob/master/LICENSE.txt
+    """
+    def __init__(self, alpha=1e-4, k=2, beta=0.75, n=5, **kwargs):
+        if n % 2 == 0:
+            raise NotImplementedError(
+                                    """LocalResponseNorm2D only works
+                                    with odd n. n provided: """ + str(n))
+        super(LocalResponseNorm2D, self).__init__(**kwargs)
+        self.alpha = alpha
+        self.k = k
+        self.beta = beta
+        self.n = n
+
+    def get_output(self, train):
+        X = self.get_input(train)
+        b, ch, r, c = K.shape(X)
+        half_n = self.n // 2
+        input_sqr = K.square(X)
+        extra_channels = K.zeros((b, ch + 2 * half_n, r, c))
+        input_sqr = K.concatenate([extra_channels[:, :half_n, :, :],
+                                   input_sqr,
+                                   extra_channels[:, half_n + ch:, :, :]],
+                                  axis=1)
+        scale = self.k
+        for i in range(self.n):
+            scale += self.alpha * input_sqr[:, i:i + ch, :, :]
+        scale = scale ** self.beta
+        return X / scale
+
+    def get_config(self):
+        config = {"name": self.__class__.__name__,
+                  "alpha": self.alpha,
+                  "k": self.k,
+                  "beta": self.beta,
+                  "n": self.n}
+        base_config = super(LocalResponseNorm2D, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+class PoolHelper(Layer):
+
+    def __init__(self, **kwargs):
+        super(PoolHelper, self).__init__(**kwargs)
+
+    def call(self, x, mask=None):
+        return x[:, :, 1:, 1:]
+
+    def get_config(self):
+        config = {}
+        base_config = super(PoolHelper, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))

yolo3/model.py

@@ -18,6 +18,9 @@
 
 from yolo3.utils import compose, update_path
 
+from yolo3.local_response_norm import LocalResponseNorm2D
+
+use_lrn = False
 
 @wraps(Conv2D)
 def DarknetConv2D(*args, **kwargs):
@@ -32,116 +35,126 @@ def DarknetConv2D_BN_Leaky(*args, **kwargs):
     """Darknet Convolution2D followed by BatchNormalization and LeakyReLU."""
     no_bias_kwargs = {'use_bias': False}
     no_bias_kwargs.update(kwargs)
+    lrn = no_bias_kwargs['lrn']
+    del no_bias_kwargs['lrn']
+    if lrn:
+        return compose(
+            DarknetConv2D(*args, **no_bias_kwargs),
+            LocalResponseNorm2D(),
+            LeakyReLU(alpha=0.1),
+        )
     return compose(
         DarknetConv2D(*args, **no_bias_kwargs),
         BatchNormalization(),
         LeakyReLU(alpha=0.1),
     )
 
 
-def resblock_body(x, num_filters, num_blocks):
+def resblock_body(x, num_filters, num_blocks, lrn):
     """A series of resblocks starting with a downsampling Convolution2D"""
     # Darknet uses left and top padding instead of 'same' mode
     x = ZeroPadding2D(((1, 0), (1, 0)))(x)
-    x = DarknetConv2D_BN_Leaky(num_filters, (3, 3), strides=(2, 2))(x)
+    x = DarknetConv2D_BN_Leaky(num_filters, (3, 3), strides=(2, 2), lrn=lrn)(x)
     for i in range(num_blocks):
         y = compose(
-            DarknetConv2D_BN_Leaky(num_filters // 2, (1, 1)),
-            DarknetConv2D_BN_Leaky(num_filters, (3, 3)))(x)
+            DarknetConv2D_BN_Leaky(num_filters // 2, (1, 1), lrn=lrn),
+            DarknetConv2D_BN_Leaky(num_filters, (3, 3), lrn=lrn))(x)
         x = Add()([x, y])
     return x
 
 
-def darknet_body(x):
+def darknet_body(x, lrn):
     """Darknent body having 52 Convolution2D layers"""
-    x = DarknetConv2D_BN_Leaky(32, (3, 3))(x)
-    x = resblock_body(x, 64, 1)
-    x = resblock_body(x, 128, 2)
-    x = resblock_body(x, 256, 8)
-    x = resblock_body(x, 512, 8)
-    x = resblock_body(x, 1024, 4)
+    x = DarknetConv2D_BN_Leaky(32, (3, 3), lrn=lrn)(x)
+    x = resblock_body(x, 64, 1, lrn=lrn)
+    x = resblock_body(x, 128, 2, lrn=lrn)
+    x = resblock_body(x, 256, 8, lrn=lrn)
+    x = resblock_body(x, 512, 8, lrn=lrn)
+    x = resblock_body(x, 1024, 4, lrn=lrn)
     return x
 
 
-def make_last_layers(x, num_filters, out_filters):
+def make_last_layers(x, num_filters, out_filters, lrn):
     """6 Conv2D_BN_Leaky layers followed by a Conv2D_linear layer"""
     x = compose(
-        DarknetConv2D_BN_Leaky(num_filters, (1, 1)),
-        DarknetConv2D_BN_Leaky(num_filters * 2, (3, 3)),
-        DarknetConv2D_BN_Leaky(num_filters, (1, 1)),
-        DarknetConv2D_BN_Leaky(num_filters * 2, (3, 3)),
-        DarknetConv2D_BN_Leaky(num_filters, (1, 1)))(x)
+        DarknetConv2D_BN_Leaky(num_filters, (1, 1), lrn=lrn),
+        DarknetConv2D_BN_Leaky(num_filters * 2, (3, 3), lrn=lrn),
+        DarknetConv2D_BN_Leaky(num_filters, (1, 1), lrn=lrn),
+        DarknetConv2D_BN_Leaky(num_filters * 2, (3, 3), lrn=lrn),
+        DarknetConv2D_BN_Leaky(num_filters, (1, 1), lrn=lrn))(x)
     y = compose(
-        DarknetConv2D_BN_Leaky(num_filters * 2, (3, 3)),
+        DarknetConv2D_BN_Leaky(num_filters * 2, (3, 3), lrn=lrn),
         DarknetConv2D(out_filters, (1, 1)))(x)
     return x, y
 
 
-def yolo_body_full(inputs, num_anchors, num_classes):
+def yolo_body_full(inputs, num_anchors, num_classes, lrn):
     """Create YOLO_V3 model CNN body in Keras.
 
     :param inputs:
     :param int num_anchors:
     :param int num_classes:
+    :param bool lrn:
     :return:
 
-    >>> yolo_body_full(Input(shape=(None, None, 3)), 6, 10)  #doctest: +ELLIPSIS
+    >>> yolo_body_full(Input(shape=(None, None, 3)), 6, 10, True)  #doctest: +ELLIPSIS
     <keras.engine.training.Model object at ...>
     """
-    darknet = Model(inputs, darknet_body(inputs))
-    x, y1 = make_last_layers(darknet.output, 512, num_anchors * (num_classes + 5))
+    darknet = Model(inputs, darknet_body(inputs, lrn=lrn))
+    x, y1 = make_last_layers(darknet.output, 512, num_anchors * (num_classes + 5), lrn=lrn)
 
     x = compose(
-        DarknetConv2D_BN_Leaky(256, (1, 1)),
+        DarknetConv2D_BN_Leaky(256, (1, 1), lrn=lrn),
         UpSampling2D(2))(x)
     x = Concatenate()([x, darknet.layers[152].output])
-    x, y2 = make_last_layers(x, 256, num_anchors * (num_classes + 5))
+    x, y2 = make_last_layers(x, 256, num_anchors * (num_classes + 5), lrn=lrn)
 
     x = compose(
-        DarknetConv2D_BN_Leaky(128, (1, 1)),
+        DarknetConv2D_BN_Leaky(128, (1, 1), lrn=lrn),
         UpSampling2D(2))(x)
     x = Concatenate()([x, darknet.layers[92].output])
-    x, y3 = make_last_layers(x, 128, num_anchors * (num_classes + 5))
+    x, y3 = make_last_layers(x, 128, num_anchors * (num_classes + 5), lrn=lrn)
 
     return Model(inputs, [y1, y2, y3])
 
 
-def yolo_body_tiny(inputs, num_anchors, num_classes):
+def yolo_body_tiny(inputs, num_anchors, num_classes, lrn):
     """Create Tiny YOLO_v3 model CNN body in keras.
 
     :param inputs:
     :param int num_anchors:
     :param int num_classes:
+    :param bool lrn:
     :return:
 
-    >>> yolo_body_tiny(Input(shape=(None, None, 3)), 6, 10)  #doctest: +ELLIPSIS
+    >>> yolo_body_tiny(Input(shape=(None, None, 3)), 6, 10, True)  #doctest: +ELLIPSIS
     <keras.engine.training.Model object at ...>
     """
     x1 = compose(
-        DarknetConv2D_BN_Leaky(16, (3, 3)),
+        DarknetConv2D_BN_Leaky(16, (3, 3), lrn=lrn),
         MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same'),
-        DarknetConv2D_BN_Leaky(32, (3, 3)),
+        DarknetConv2D_BN_Leaky(32, (3, 3), lrn=lrn),
         MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same'),
-        DarknetConv2D_BN_Leaky(64, (3, 3)),
+        DarknetConv2D_BN_Leaky(64, (3, 3), lrn=lrn),
         MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same'),
-        DarknetConv2D_BN_Leaky(128, (3, 3)),
+        DarknetConv2D_BN_Leaky(128, (3, 3), lrn=lrn),
         MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same'),
-        DarknetConv2D_BN_Leaky(256, (3, 3)))(inputs)
+        DarknetConv2D_BN_Leaky(256, (3, 3), lrn=lrn))(inputs)
     x2 = compose(
         MaxPooling2D(pool_size=(2, 2), strides=(2, 2), padding='same'),
-        DarknetConv2D_BN_Leaky(512, (3, 3)),
+        DarknetConv2D_BN_Leaky(512, (3, 3), lrn=lrn),
         MaxPooling2D(pool_size=(2, 2), strides=(1, 1), padding='same'),
-        DarknetConv2D_BN_Leaky(1024, (3, 3)),
-        DarknetConv2D_BN_Leaky(256, (1, 1)))(x1)
+        DarknetConv2D_BN_Leaky(1024, (3, 3), lrn=lrn),
+        DarknetConv2D_BN_Leaky(256, (1, 1), lrn=lrn))(x1)
     y1 = compose(
-        DarknetConv2D_BN_Leaky(512, (3, 3)),
+        DarknetConv2D_BN_Leaky(512, (3, 3), lrn=lrn),
         DarknetConv2D(num_anchors * (num_classes + 5), (1, 1)))(x2)
     x3 = compose(
-        DarknetConv2D_BN_Leaky(128, (1, 1)),
+        DarknetConv2D_BN_Leaky(128, (1, 1), lrn=lrn),
         UpSampling2D(2))(x2)
     y2 = compose(
         Concatenate(),
-        DarknetConv2D_BN_Leaky(256, (3, 3)),
+        DarknetConv2D_BN_Leaky(256, (3, 3), lrn=lrn),
         DarknetConv2D(num_anchors * (num_classes + 5), (1, 1)))([x3, x1])
 
     return Model(inputs, [y1, y2])
@@ -399,7 +412,7 @@ def _loop_body(b, ignore_mask):
 
 
 def create_model(input_shape, anchors, num_classes, weights_path=None, model_factor=3,
-                 freeze_body=2, ignore_thresh=0.5, nb_gpu=1):
+                 freeze_body=2, ignore_thresh=0.5, nb_gpu=1, lrn=False):
     """create the training model"""
     _INPUT_SHAPES = {0: 32, 1: 16, 2: 8, 3: 4}
     _FACTOR_YOLO_BODY = {2: yolo_body_tiny, 3: yolo_body_full}
@@ -417,7 +430,7 @@ def create_model(input_shape, anchors, num_classes, weights_path=None, model_fac
     h, w = input_shape
     num_anchors = len(anchors)
 
-    model_body = _FACTOR_YOLO_BODY[model_factor](image_input, num_anchors // model_factor, num_classes)
+    model_body = _FACTOR_YOLO_BODY[model_factor](image_input, num_anchors // model_factor, num_classes, lrn=lrn)
     logging.debug('Create YOLOv3 (model_factor: %i) model with %i anchors and %i classes.',
                   model_factor, num_anchors, num_classes)
 
@@ -453,15 +466,15 @@ def create_model(input_shape, anchors, num_classes, weights_path=None, model_fac
 
 
 def create_model_tiny(input_shape, anchors, num_classes, weights_path=None,
-                      freeze_body=2, ignore_thresh=0.5, nb_gpu=1):
+                      freeze_body=2, ignore_thresh=0.5, nb_gpu=1, lrn=False):
     """create the training model, for Tiny YOLOv3 """
 
     return create_model(input_shape, anchors, num_classes, weights_path, model_factor=2,
-                        freeze_body=freeze_body, ignore_thresh=ignore_thresh, nb_gpu=nb_gpu)
+                        freeze_body=freeze_body, ignore_thresh=ignore_thresh, nb_gpu=nb_gpu, lrn=lrn)
 
 
 def create_model_bottleneck(input_shape, anchors, num_classes, freeze_body=2,
-                            weights_path=None, nb_gpu=1):
+                            weights_path=None, nb_gpu=1, lrn=False):
     """create the training model"""
     # K.clear_session()  # get a new session
     image_input = Input(shape=(None, None, 3))
@@ -482,7 +495,7 @@ def create_model_bottleneck(input_shape, anchors, num_classes, freeze_body=2,
     if not nb_gpu:  # disable all GPUs
         os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 
-    model_body = yolo_body_full(image_input, num_anchors // 3, num_classes)
+    model_body = yolo_body_full(image_input, num_anchors // 3, num_classes, lrn=lrn)
     logging.info('Create YOLOv3 model with %i anchors and %i classes.',
                  num_anchors, num_classes)