fix PredictConfig. use inputs_desc adn tower_func for load-cpm

examples/ConvolutionalPoseMachines/load-cpm.py

@@ -44,71 +44,65 @@ def get_gaussian_map():
     return gaussian_map.reshape((1, 368, 368, 1))
 
 
-class Model(ModelDesc):
-    def _get_inputs(self):
-        return [InputDesc(tf.float32, (None, 368, 368, 3), 'input'),
-                InputDesc(tf.float32, (None, 368, 368, 15), 'label'),
-                ]
-
-    def _build_graph(self, inputs):
-        image, label = inputs
-        image = image / 256.0 - 0.5
-
-        gmap = tf.constant(get_gaussian_map())
-        gmap = tf.pad(gmap, [[0, 0], [0, 1], [0, 1], [0, 0]])
-        pool_center = AvgPooling('mappool', gmap, 9, stride=8, padding='VALID')
-        with argscope(Conv2D, kernel_shape=3, nl=tf.nn.relu,
-                      W_init=tf.random_normal_initializer(stddev=0.01)):
-            shared = (LinearWrap(image)
-                      .Conv2D('conv1_1', 64)
-                      .Conv2D('conv1_2', 64)
-                      .MaxPooling('pool1', 2)
-                      # 184
-                      .Conv2D('conv2_1', 128)
-                      .Conv2D('conv2_2', 128)
-                      .MaxPooling('pool2', 2)
-                      # 92
-                      .Conv2D('conv3_1', 256)
-                      .Conv2D('conv3_2', 256)
-                      .Conv2D('conv3_3', 256)
-                      .Conv2D('conv3_4', 256)
-                      .MaxPooling('pool3', 2)
-                      # 46
-                      .Conv2D('conv4_1', 512)
-                      .Conv2D('conv4_2', 512)
-                      .Conv2D('conv4_3_CPM', 256)
-                      .Conv2D('conv4_4_CPM', 256)
-                      .Conv2D('conv4_5_CPM', 256)
-                      .Conv2D('conv4_6_CPM', 256)
-                      .Conv2D('conv4_7_CPM', 128)())
-
-        def add_stage(stage, l):
-            l = tf.concat([l, shared, pool_center], 3,
-                          name='concat_stage{}'.format(stage))
-            for i in range(1, 6):
-                l = Conv2D('Mconv{}_stage{}'.format(i, stage), l, 128)
-            l = Conv2D('Mconv6_stage{}'.format(stage), l, 128, kernel_shape=1)
-            l = Conv2D('Mconv7_stage{}'.format(stage),
-                       l, 15, kernel_shape=1, nl=tf.identity)
-            return l
-
-        with argscope(Conv2D, kernel_shape=7, nl=tf.nn.relu):
-            out1 = (LinearWrap(shared)
-                    .Conv2D('conv5_1_CPM', 512, kernel_shape=1)
-                    .Conv2D('conv5_2_CPM', 15, kernel_shape=1, nl=tf.identity)())
-            out2 = add_stage(2, out1)
-            out3 = add_stage(3, out2)
-            out4 = add_stage(4, out3)
-            out5 = add_stage(5, out4)
-            out6 = add_stage(6, out4)
-            resized_map = tf.image.resize_bilinear(out6,
-                                                   [368, 368], name='resized_map')
+def CPM(image):
+    image = image / 256.0 - 0.5
+
+    gmap = tf.constant(get_gaussian_map())
+    gmap = tf.pad(gmap, [[0, 0], [0, 1], [0, 1], [0, 0]])
+    pool_center = AvgPooling('mappool', gmap, 9, stride=8, padding='VALID')
+    with argscope(Conv2D, kernel_shape=3, nl=tf.nn.relu,
+                  W_init=tf.random_normal_initializer(stddev=0.01)):
+        shared = (LinearWrap(image)
+                  .Conv2D('conv1_1', 64)
+                  .Conv2D('conv1_2', 64)
+                  .MaxPooling('pool1', 2)
+                  # 184
+                  .Conv2D('conv2_1', 128)
+                  .Conv2D('conv2_2', 128)
+                  .MaxPooling('pool2', 2)
+                  # 92
+                  .Conv2D('conv3_1', 256)
+                  .Conv2D('conv3_2', 256)
+                  .Conv2D('conv3_3', 256)
+                  .Conv2D('conv3_4', 256)
+                  .MaxPooling('pool3', 2)
+                  # 46
+                  .Conv2D('conv4_1', 512)
+                  .Conv2D('conv4_2', 512)
+                  .Conv2D('conv4_3_CPM', 256)
+                  .Conv2D('conv4_4_CPM', 256)
+                  .Conv2D('conv4_5_CPM', 256)
+                  .Conv2D('conv4_6_CPM', 256)
+                  .Conv2D('conv4_7_CPM', 128)())
+
+    def add_stage(stage, l):
+        l = tf.concat([l, shared, pool_center], 3,
+                      name='concat_stage{}'.format(stage))
+        for i in range(1, 6):
+            l = Conv2D('Mconv{}_stage{}'.format(i, stage), l, 128)
+        l = Conv2D('Mconv6_stage{}'.format(stage), l, 128, kernel_shape=1)
+        l = Conv2D('Mconv7_stage{}'.format(stage),
+                   l, 15, kernel_shape=1, nl=tf.identity)
+        return l
+
+    with argscope(Conv2D, kernel_shape=7, nl=tf.nn.relu):
+        out1 = (LinearWrap(shared)
+                .Conv2D('conv5_1_CPM', 512, kernel_shape=1)
+                .Conv2D('conv5_2_CPM', 15, kernel_shape=1, nl=tf.identity)())
+        out2 = add_stage(2, out1)
+        out3 = add_stage(3, out2)
+        out4 = add_stage(4, out3)
+        out5 = add_stage(5, out4)
+        out6 = add_stage(6, out4)
+        resized_map = tf.image.resize_bilinear(out6,
+                                               [368, 368], name='resized_map')
 
 
 def run_test(model_path, img_file):
     param_dict = np.load(model_path, encoding='latin1').item()
     predict_func = OfflinePredictor(PredictConfig(
-        model=Model(),
+        inputs_desc=[InputDesc(tf.float32, (None, 368, 368, 3), 'input')],
+        tower_func=CPM,
         session_init=DictRestore(param_dict),
         input_names=['input'],
         output_names=['resized_map']

tensorpack/predict/config.py

@@ -27,7 +27,7 @@ class PredictConfig(object):
                  ):
         """
         Args:
-            model (ModelDescBase): the model to be used to obtain inputs_desc and tower_func.
+            model (ModelDescBase): to be used to obtain inputs_desc and tower_func.
             inputs_desc ([InputDesc]):
             tower_func: a callable which takes input tensors
 
@@ -35,8 +35,7 @@ class PredictConfig(object):
                 session. Defaults to :class:`tf.train.ChiefSessionCreator()`.
             session_init (SessionInit): how to initialize variables of the session.
                 Defaults to do nothing.
-            input_names (list): a list of input tensor names. Defaults to all
-                inputs of the model.
+            input_names (list): a list of input tensor names. Defaults to match inputs_desc.
             output_names (list): a list of names of the output tensors to predict, the
                 tensors can be any computable tensor in the graph.
             return_input (bool): same as in :attr:`PredictorBase.return_input`.
@@ -70,8 +69,7 @@ class PredictConfig(object):
         # inputs & outputs
         self.input_names = input_names
         if self.input_names is None:
-            raw_tensors = self.model.get_inputs_desc()
-            self.input_names = [k.name for k in raw_tensors]
+            self.input_names = [k.name for k in self.inputs_desc]
         self.output_names = output_names
         assert_type(self.output_names, list)
         assert_type(self.input_names, list)