update model

bashes/basic_cnn.sh

@@ -2,6 +2,6 @@
 
 cd ..
 python3 train.py\
-    --model_class BasicCNNModel\
-    --checkpoint_dir checkpoints/basic_cnn\
-    --tensor_board_dir events/basic_cnn
+    --model_class_name BasicCNNModel\
+    --checkpoint_dir checkpoints/basic\
+    --tensor_board_dir events/basic

bashes/vgg19.sh

@@ -2,6 +2,6 @@
 
 cd ..
 python3 train.py\
-    --model_class VGG19Model\
+    --model_class_name VGG19Model\
     --checkpoint_dir checkpoints/vgg19\
     --tensor_board_dir events/vgg19

models/__init__.py

@@ -0,0 +1,2 @@
+from .basic_cnn import BasicCNNModel
+from .vgg19 import VGG19Model

models/basic_cnn.py

@@ -0,0 +1,68 @@
+from model_zoo.model import BaseModel
+import tensorflow as tf
+
+
+class BasicCNNModel(BaseModel):
+    """
+    Basic CNN Model
+    """
+
+    def __init__(self, config):
+        """
+        init layers
+        :param config:
+        """
+        super(BasicCNNModel, self).__init__(config)
+        self.bn1 = tf.keras.layers.BatchNormalization()
+        self.conv1 = tf.keras.layers.Conv2D(32, (2, 2), padding='same', activation='relu',
+                                            kernel_initializer='random_uniform')
+        self.pool1 = tf.keras.layers.MaxPool2D(padding='same')
+        self.dropout1 = tf.keras.layers.Dropout(0.5)
+        self.conv2 = tf.keras.layers.Conv2D(32, (2, 2), padding='same', activation='relu',
+                                            kernel_initializer='random_uniform')
+        self.pool2 = tf.keras.layers.MaxPool2D(padding='same')
+        self.dropout2 = tf.keras.layers.Dropout(0.5)
+        self.flatten1 = tf.keras.layers.Flatten()
+        self.dense1 = tf.keras.layers.Dense(128, activation='relu', kernel_initializer='random_uniform')
+        self.dense2 = tf.keras.layers.Dense(10, activation='softmax')
+
+    def call(self, inputs, training=False, mask=None):
+        """
+        build model
+        :param inputs: inputs image
+        :param training:
+        :param mask:
+        :return:
+        """
+        o = self.bn1(inputs)
+        o = self.conv1(o)
+        o = self.pool1(o)
+        o = self.dropout1(o) if training else o
+        o = self.conv2(o)
+        o = self.pool2(o)
+        o = self.dropout2(o) if training else o
+        o = self.flatten1(o)
+        o = self.dense1(o)
+        o = self.dense2(o)
+        return o
+
+    def get_optimizer(self):
+        """
+        build optimizer
+        :return:
+        """
+        return tf.keras.optimizers.Adam(lr=self.config.get('learning_rate'))
+
+    def get_loss(self):
+        """
+        define loss
+        :return:
+        """
+        return 'categorical_crossentropy'
+
+    def get_metrics(self):
+        """
+        define metrics
+        :return:
+        """
+        return ['accuracy']

model.py → models/vgg19.py

@@ -2,46 +2,11 @@
 import tensorflow as tf
 
 
-class BasicCNNModel(BaseModel):
-    def __init__(self, config):
-        super(BasicCNNModel, self).__init__(config)
-        self.bn1 = tf.keras.layers.BatchNormalization()
-        self.conv1 = tf.keras.layers.Conv2D(32, (2, 2), padding='same', activation='relu',
-                                            kernel_initializer='random_uniform')
-        self.pool1 = tf.keras.layers.MaxPool2D(padding='same')
-        self.dropout1 = tf.keras.layers.Dropout(0.5)
-        self.conv2 = tf.keras.layers.Conv2D(32, (2, 2), padding='same', activation='relu',
-                                            kernel_initializer='random_uniform')
-        self.pool2 = tf.keras.layers.MaxPool2D(padding='same')
-        self.dropout2 = tf.keras.layers.Dropout(0.5)
-        self.flatten1 = tf.keras.layers.Flatten()
-        self.dense1 = tf.keras.layers.Dense(128, activation='relu', kernel_initializer='random_uniform')
-        self.dense2 = tf.keras.layers.Dense(10, activation='softmax')
-
-    def call(self, inputs, training=None, mask=None):
-        o = self.bn1(inputs)
-        o = self.conv1(o)
-        o = self.pool1(o)
-        o = self.dropout1(o) if training else o
-        o = self.conv2(o)
-        o = self.pool2(o)
-        o = self.dropout2(o) if training else o
-        o = self.flatten1(o)
-        o = self.dense1(o)
-        o = self.dense2(o)
-        return o
-
-    def optimizer(self):
-        return tf.train.AdamOptimizer(self.config.get('learning_rate'))
-
-    def init(self):
-        self.compile(optimizer=self.optimizer(), loss='categorical_crossentropy', metrics=['accuracy'])
-
-
 class VGG19Model(BaseModel):
     """
     This model gets no good results, deprecated.
     """
+
     def __init__(self, config):
         super(VGG19Model, self).__init__(config)
         self.num_features = 64
@@ -53,7 +18,7 @@ def __init__(self, config):
         self.bn1 = tf.keras.layers.BatchNormalization()
         self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
         self.drop1 = tf.keras.layers.Dropout(rate=0.5)
-        
+
         # layer2
         self.conv21 = tf.keras.layers.Conv2D(filters=2 * self.num_features, kernel_size=(3, 3), activation='relu',
                                              padding='same')
@@ -62,7 +27,7 @@ def __init__(self, config):
         self.bn2 = tf.keras.layers.BatchNormalization()
         self.pool2 = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
         self.drop2 = tf.keras.layers.Dropout(rate=0.5)
-        
+
         # layer3
         self.conv31 = tf.keras.layers.Conv2D(filters=2 * 2 * self.num_features, kernel_size=(3, 3), activation='relu',
                                              padding='same')
@@ -73,7 +38,7 @@ def __init__(self, config):
         self.bn3 = tf.keras.layers.BatchNormalization()
         self.pool3 = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
         self.drop3 = tf.keras.layers.Dropout(rate=0.5)
-        
+
         # layer4
         self.conv41 = tf.keras.layers.Conv2D(filters=2 * 2 * 2 * self.num_features, kernel_size=(3, 3),
                                              activation='relu',
@@ -87,7 +52,7 @@ def __init__(self, config):
         self.bn4 = tf.keras.layers.BatchNormalization()
         self.pool4 = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
         self.drop4 = tf.keras.layers.Dropout(rate=0.5)
-        
+
         # layer5
         self.conv51 = tf.keras.layers.Conv2D(filters=2 * 2 * 2 * self.num_features, kernel_size=(3, 3),
                                              activation='relu',
@@ -101,20 +66,20 @@ def __init__(self, config):
         self.bn5 = tf.keras.layers.BatchNormalization()
         self.pool5 = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))
         self.drop5 = tf.keras.layers.Dropout(rate=0.5)
-        
+
         # flatten
         self.flatten = tf.keras.layers.Flatten()
-        
+
         # dense
         self.dense1 = tf.keras.layers.Dense(2 * 2 * 2 * self.num_features, activation='relu')
         self.drop5 = tf.keras.layers.Dropout(0.5)
         self.dense2 = tf.keras.layers.Dense(2 * 2 * self.num_features, activation='relu')
         self.drop6 = tf.keras.layers.Dropout(0.5)
         self.dense3 = tf.keras.layers.Dense(2 * self.num_features, activation='relu')
         self.drop7 = tf.keras.layers.Dropout(0.5)
-        
+
         self.dense4 = tf.keras.layers.Dense(10, activation='softmax')
-    
+
     def call(self, inputs, training=None, mask=None):
         # layer1
         x = self.conv11(inputs)
@@ -149,7 +114,7 @@ def call(self, inputs, training=None, mask=None):
         x = self.bn5(x, training=training)
         x = self.pool5(x)
         x = self.drop5(x, training=training)
-        
+
         # flatten
         x = self.flatten(x)
         # dense
@@ -161,9 +126,24 @@ def call(self, inputs, training=None, mask=None):
         x = self.drop7(x, training=training)
         x = self.dense4(x)
         return x
-
-    def optimizer(self):
-        return tf.train.AdamOptimizer(self.config.get('learning_rate'))
-
-    def init(self):
-        self.compile(optimizer=self.optimizer(), loss='categorical_crossentropy', metrics=['accuracy'])
+
+    def get_optimizer(self):
+        """
+        build optimizer
+        :return:
+        """
+        return tf.keras.optimizers.Adam(lr=self.config.get('learning_rate'))
+
+    def get_loss(self):
+        """
+        define loss
+        :return:
+        """
+        return 'categorical_crossentropy'
+
+    def get_metrics(self):
+        """
+        define metrics
+        :return:
+        """
+        return ['accuracy']

train.py

@@ -1,17 +1,25 @@
 import tensorflow as tf
 from model_zoo.trainer import BaseTrainer
-from tensorflow.python.keras.datasets import fashion_mnist
+from model_zoo import datasets
+from model_zoo import flags
 
-tf.flags.DEFINE_integer('epochs', 100, 'Max epochs')
-tf.flags.DEFINE_float('learning_rate', 0.01, 'Learning rate')
-tf.flags.DEFINE_string('model_class', 'VGG19Model', help='Model class name')
-tf.flags.DEFINE_string('checkpoint_name', 'model.ckpt', help='Model name')
+flags.DEFINE_integer('epochs', 100)
+flags.DEFINE_float('learning_rate', 0.01)
+flags.DEFINE_string('model_class_name', 'BasicCNNModel')
+flags.DEFINE_string('checkpoint_name', 'model.ckpt')
 
 
 class Trainer(BaseTrainer):
-
+    """
+    Train Image Classification Model.
+    """
+
     def prepare_data(self):
-        (x_train, y_train), (_, _) = fashion_mnist.load_data()
+        """
+        Prepare fashion mnist data.
+        :return:
+        """
+        (x_train, y_train), (_, _) = datasets.fashion_mnist.load_data()
         x_train = x_train.reshape((-1, 28, 28, 1))
         x_train, y_train = x_train.astype('float16') / 255.0, \
                            tf.keras.utils.to_categorical(y_train.astype('float16'), 10)