update ModelHelper's func interface in the examplar code

jiaxiang-wu · jiaxiang-wu · commit 9423e161e8f8 · 2019-03-03T16:23:23.000+08:00
diff --git a/examples/convnet_at_fmnist.py b/examples/convnet_at_fmnist.py
@@ -69,11 +69,11 @@ def forward_fn(inputs, data_format):
 class ModelHelper(AbstractModelHelper):
   """Model helper for creating a ConvNet model for the Fashion-MNIST dataset."""
 
-  def __init__(self):
+  def __init__(self, data_format='channels_last'):
     """Constructor function."""
 
     # class-independent initialization
-    super(ModelHelper, self).__init__()
+    super(ModelHelper, self).__init__(data_format)
 
     # initialize training & evaluation subsets
     self.dataset_train = FMnistDataset(is_train=True)
@@ -89,15 +89,15 @@ def build_dataset_eval(self):
 
     return self.dataset_eval.build()
 
-  def forward_train(self, inputs, data_format='channels_last'):
+  def forward_train(self, inputs):
     """Forward computation at training."""
 
-    return forward_fn(inputs, data_format)
+    return forward_fn(inputs, self.data_format)
 
-  def forward_eval(self, inputs, data_format='channels_last'):
+  def forward_eval(self, inputs):
     """Forward computation at evaluation."""
 
-    return forward_fn(inputs, data_format)
+    return forward_fn(inputs, self.data_format)
 
   def calc_loss(self, labels, outputs, trainable_vars):
     """Calculate loss (and some extra evaluation metrics)."""
