fix BN axis error

Author: chaoyan1037

DenseNet/densenet.py

@@ -9,10 +9,12 @@
 import keras.backend as K
 
 
-def conv_factory(x, nb_filter, dropout_rate=None, weight_decay=1E-4):
+def conv_factory(x, concat_axis, nb_filter,
+                 dropout_rate=None, weight_decay=1E-4):
     """Apply BatchNorm, Relu 3x3Conv2D, optional dropout
 
     :param x: Input keras network
+    :param concat_axis: int -- index of contatenate axis
     :param nb_filter: int -- number of filters
     :param dropout_rate: int -- dropout rate
     :param weight_decay: int -- weight decay factor
@@ -21,7 +23,7 @@ def conv_factory(x, nb_filter, dropout_rate=None, weight_decay=1E-4):
     :rtype: keras network
     """
 
-    x = BatchNormalization(axis=1,
+    x = BatchNormalization(axis=concat_axis,
                            gamma_regularizer=l2(weight_decay),
                            beta_regularizer=l2(weight_decay))(x)
     x = Activation('relu')(x)
@@ -36,10 +38,12 @@ def conv_factory(x, nb_filter, dropout_rate=None, weight_decay=1E-4):
     return x
 
 
-def transition(x, nb_filter, dropout_rate=None, weight_decay=1E-4):
+def transition(x, concat_axis, nb_filter,
+               dropout_rate=None, weight_decay=1E-4):
     """Apply BatchNorm, Relu 1x1Conv2D, optional dropout and Maxpooling2D
 
     :param x: keras model
+    :param concat_axis: int -- index of contatenate axis
     :param nb_filter: int -- number of filters
     :param dropout_rate: int -- dropout rate
     :param weight_decay: int -- weight decay factor
@@ -49,7 +53,7 @@ def transition(x, nb_filter, dropout_rate=None, weight_decay=1E-4):
 
     """
 
-    x = BatchNormalization(axis=1,
+    x = BatchNormalization(axis=concat_axis,
                            gamma_regularizer=l2(weight_decay),
                            beta_regularizer=l2(weight_decay))(x)
     x = Activation('relu')(x)
@@ -65,12 +69,13 @@ def transition(x, nb_filter, dropout_rate=None, weight_decay=1E-4):
     return x
 
 
-def denseblock(x, nb_layers, nb_filter, growth_rate,
+def denseblock(x, concat_axis, nb_layers, nb_filter, growth_rate,
                dropout_rate=None, weight_decay=1E-4):
     """Build a denseblock where the output of each
        conv_factory is fed to subsequent ones
 
     :param x: keras model
+    :param concat_axis: int -- index of contatenate axis
     :param nb_layers: int -- the number of layers of conv_
                       factory to append to the model.
     :param nb_filter: int -- number of filters
@@ -84,26 +89,23 @@ def denseblock(x, nb_layers, nb_filter, growth_rate,
 
     list_feat = [x]
 
-    if K.image_dim_ordering() == "th":
-        concat_axis = 1
-    elif K.image_dim_ordering() == "tf":
-        concat_axis = -1
-
     for i in range(nb_layers):
-        x = conv_factory(x, growth_rate, dropout_rate, weight_decay)
+        x = conv_factory(x, concat_axis, growth_rate,
+                         dropout_rate, weight_decay)
         list_feat.append(x)
         x = Concatenate(axis=concat_axis)(list_feat)
         nb_filter += growth_rate
 
     return x, nb_filter
 
 
-def denseblock_altern(x, nb_layers, nb_filter, growth_rate,
+def denseblock_altern(x, concat_axis, nb_layers, nb_filter, growth_rate,
                       dropout_rate=None, weight_decay=1E-4):
     """Build a denseblock where the output of each conv_factory
        is fed to subsequent ones. (Alternative of a above)
 
     :param x: keras model
+    :param concat_axis: int -- index of contatenate axis
     :param nb_layers: int -- the number of layers of conv_
                       factory to append to the model.
     :param nb_filter: int -- number of filters
@@ -117,13 +119,9 @@ def denseblock_altern(x, nb_layers, nb_filter, growth_rate,
     above is that the one above
     """
 
-    if K.image_dim_ordering() == "th":
-        concat_axis = 1
-    elif K.image_dim_ordering() == "tf":
-        concat_axis = -1
-
     for i in range(nb_layers):
-        merge_tensor = conv_factory(x, growth_rate, dropout_rate, weight_decay)
+        merge_tensor = conv_factory(x, concat_axis, growth_rate,
+                                    dropout_rate, weight_decay)
         x = Concatenate(axis=concat_axis)([merge_tensor, x])
         nb_filter += growth_rate
 
@@ -147,6 +145,11 @@ def DenseNet(nb_classes, img_dim, depth, nb_dense_block, growth_rate,
     :rtype: keras model
 
     """
+    
+    if K.image_dim_ordering() == "th":
+        concat_axis = 1
+    elif K.image_dim_ordering() == "tf":
+        concat_axis = -1
 
     model_input = Input(shape=img_dim)
 
@@ -165,19 +168,21 @@ def DenseNet(nb_classes, img_dim, depth, nb_dense_block, growth_rate,
 
     # Add dense blocks
     for block_idx in range(nb_dense_block - 1):
-        x, nb_filter = denseblock(x, nb_layers, nb_filter, growth_rate,
+        x, nb_filter = denseblock(x, concat_axis, nb_layers,
+                                  nb_filter, growth_rate, 
                                   dropout_rate=dropout_rate,
                                   weight_decay=weight_decay)
         # add transition
         x = transition(x, nb_filter, dropout_rate=dropout_rate,
                        weight_decay=weight_decay)
 
     # The last denseblock does not have a transition
-    x, nb_filter = denseblock(x, nb_layers, nb_filter, growth_rate,
+    x, nb_filter = denseblock(x, concat_axis, nb_layers,
+                              nb_filter, growth_rate, 
                               dropout_rate=dropout_rate,
                               weight_decay=weight_decay)
 
-    x = BatchNormalization(axis=1,
+    x = BatchNormalization(axis=concat_axis,
                            gamma_regularizer=l2(weight_decay),
                            beta_regularizer=l2(weight_decay))(x)
     x = Activation('relu')(x)