Refactor mnist convnet for multi GPU usage

Author: blazekadam

mnist_convnet/config.yaml

@@ -19,5 +19,7 @@ hooks:
   - ComputeStats:
       variables: [loss, accuracy]
   - LogVariables
+  - LogProfile
   - StopAfter:
       minutes: 2
+

mnist_convnet/convnet.py

@@ -6,20 +6,19 @@
 class SimpleConvNet(cxtf.BaseModel):
 
     def _create_model(self):
-        images = tf.placeholder(tf.float32, shape=[None, 28, 28], name='images')
+        images = tf.placeholder(tf.float32, shape=[None, 28, 28, 1], name='images')
         labels = tf.placeholder(tf.int64, shape=[None], name='labels')
 
         with tf.variable_scope('conv1'):
-            net = tf.expand_dims(images, -1)
-            net = K.layers.Conv2D(20, 5)(net)
+            net = K.layers.Conv2D(64, 5)(images)
             net = K.layers.MaxPool2D()(net)
         with tf.variable_scope('conv2'):
-            net = K.layers.Conv2D(50, 3)(net)
+            net = K.layers.Conv2D(128, 3)(net)
             net = K.layers.MaxPool2D()(net)
         with tf.variable_scope('dense3'):
             net = K.layers.Flatten()(net)
             net = K.layers.Dropout(0.4).apply(net, training=self.is_training)
-            net = K.layers.Dense(100)(net)
+            net = K.layers.Dense(64)(net)
         with tf.variable_scope('dense4'):
             logits = K.layers.Dense(10, activation=None)(net)
 

mnist_convnet/mnist_dataset.py

@@ -35,7 +35,7 @@ def _load_data(self) -> None:
                 with gzip.open(file_path, 'rb') as file:
                     if 'images' in key:
                         _, _, rows, cols = struct.unpack(">IIII", file.read(16))
-                        self._data[key] = np.frombuffer(file.read(), dtype=np.uint8).reshape(-1, rows, cols)
+                        self._data[key] = np.frombuffer(file.read(), dtype=np.uint8).reshape(-1, rows, cols, 1)
                     else:
                         _ = struct.unpack(">II", file.read(8))
                         self._data[key] = np.frombuffer(file.read(), dtype=np.int8)