Added use_batch_norm param

ZFTurbo · May 22, 2022 · ccc89ff · ccc89ff
1 parent 267eb8c
commit ccc89ff
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/classification_models_1D/models/vgg16.py b/classification_models_1D/models/vgg16.py
@@ -11,6 +11,7 @@
 from __future__ import print_function
 
 import os
+import numpy as np
 from .. import get_submodules_from_kwargs
 from ..weights import load_model_weights
 from keras_applications import imagenet_utils
@@ -32,6 +33,7 @@ def VGG16(
         init_filters=64,
         max_filters=512,
         repetitions=(2, 2, 3, 3, 3),
+        use_batch_norm=False,
         **kwargs
 ):
     """Instantiates the VGG16 architecture.
@@ -121,6 +123,8 @@ def VGG16(
                 padding='same',
                 name='block{}_conv{}'.format(stage + 1, i + 1)
             )(x)
+            if use_batch_norm:
+                x = layers.BatchNormalization(epsilon=1e-5, momentum=0.1)(x)
 
         x = layers.MaxPooling1D(stride_size[stage], strides=stride_size[stage], name='block{}_pool'.format(stage+1))(x)
 

diff --git a/classification_models_1D/models/vgg19.py b/classification_models_1D/models/vgg19.py
@@ -32,6 +32,7 @@ def VGG19(
         init_filters=64,
         max_filters=512,
         repetitions=(2, 2, 4, 4, 4),
+        use_batch_norm=False,
         **kwargs
 ):
     """Instantiates the VGG19 architecture.
@@ -121,6 +122,8 @@ def VGG19(
                 padding='same',
                 name='block{}_conv{}'.format(stage + 1, i + 1)
             )(x)
+            if use_batch_norm:
+                x = layers.BatchNormalization(epsilon=1e-5, momentum=0.1)(x)
 
         x = layers.MaxPooling1D(stride_size[stage], strides=stride_size[stage], name='block{}_pool'.format(stage + 1))(
             x)