Add the most simple training example to show-case training

2026-09-12 12:01:02 +08:00 · 2018-11-06 14:20:18 +01:00
parent 2cbeb7e699
commit 4ec3734736
7 changed files with 344 additions and 93 deletions
@@ -0,0 +1,268 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# Copyright 2017 ComputerGraphics Tuebingen. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+
+
+"""
+Run 3D-MNIST classification using Flex-Convolutions, without any fancy parts.
+
+Example-output:
+
+ python mnist_3d.py --gpu 0 -fusion pooling
+
+ [@base.py:282] Epoch 56 (global_step 210000) finished, time:23 minutes 52 seconds.
+ [@saver.py:77] Model saved to train_log/conv_position_pooling/model-210000.
+ 100%|##############################|625/625[00:55<00:00,11.28it/s]
+ [@monitor.py:459] accuracy: 0.94699
+ [@monitor.py:459] learning_rate: 0.0001
+ [@monitor.py:459] train_error: 0.053013
+ [@monitor.py:459] validation_accuracy: 0.9327
+ [@monitor.py:459] validation_cross_entropy_loss: 0.209
+ [@group.py:48] Callbacks took 55.513 sec in total. InferenceRunner: 55.4 seconds
+ [@base.py:272] Start Epoch 57 ...
+
+
+ python mnist_3d.py --gpu 0 -fusion conv
+
+ [@base.py:282] Epoch 40 (global_step 150000) finished, time:34 minutes 49 seconds.
+ [@saver.py:77] Model saved to train_log/conv_position_conv/model-150000.
+ 100%|##############################|625/625[00:57<00:00,10.89it/s]
+ [@monitor.py:459] accuracy: 0.89019
+ [@monitor.py:459] learning_rate: 0.0001
+ [@monitor.py:459] train_error: 0.10981
+ [@monitor.py:459] validation_accuracy: 0.8748
+ [@monitor.py:459] validation_cross_entropy_loss: 0.38396
+
+This implementation is based on Tensorpack
+- http://tensorpack.com/
+- https://github.com/tensorpack/tensorpack/
+
+"""
+
+import os
+import argparse
+import tensorflow as tf
+import numpy as np
+import cv2
+
+from tensorpack import *
+from layers import flex_convolution, flex_pooling, knn_bruteforce
+
+
+enable_argscope_for_module(tf.layers)
+
+TOTAL_BATCH_SIZE = 16
+BATCH_SIZE = 16
+SHAPE = 28
+CHANNELS = 3
+
+USE_POOLING = False
+
+PC = {'num': 1024, 'dp': 3}
+
+
+class Digit2Cloud(RNGDataFlow):
+  """ A very basic 2D-MNIST to 3D-MNIST sampler on a regular grid.
+  """
+  def __init__(self, incoming_df, num=1024):
+    super(Digit2Cloud, self).__init__()
+    self.incoming_df = incoming_df
+    self.num = num
+
+  def reset_state(self):
+    super(Digit2Cloud, self).reset_state()
+    self.incoming_df.reset_state()
+
+  def __len__(self):
+    return self.incoming_df.__len__()
+
+  def map(self, dp, num=1024):
+    digit = dp[0]
+
+    # detect edges
+    def auto_canny(image, sigma=0.33):
+      v = np.median(image)
+      lower = int(max(0, (1.0 - sigma) * v))
+      upper = int(min(255, (1.0 + sigma) * v))
+      edged = cv2.Canny(image, lower, upper)
+      return edged
+
+    digit = np.tile(np.expand_dims(digit, axis=-1), [1, 1, 3])
+    digit = cv2.resize(digit, (32, 32))
+
+    img = (255 * digit).astype(np.uint8)
+    blurred = cv2.GaussianBlur(img, (3, 3), 0)
+    canny = auto_canny(blurred)
+    edge_x, edge_y = np.nonzero(canny > 0)
+
+    xlen = np.max(edge_x) - np.min(edge_x)
+    ylen = np.max(edge_y) - np.min(edge_y)
+
+    face_x, face_y = np.nonzero(img[:, :, 0] > 0)
+    points = []
+
+    depth = max(ylen, xlen)
+    padding = (32 - depth) / 2.
+
+    def z_dim(x):
+      return x * depth + padding
+
+    # start sampling (just extrude digits)
+    for i in range(1024):
+      choice = self.rng.randint(2 + 4)
+      if choice > 1:
+        idx = self.rng.randint(len(edge_x))
+        z = self.rng.rand()
+        points.append([edge_x[idx], edge_y[idx], z_dim(z)])
+      else:
+        idx = self.rng.randint(len(face_x))
+        z = self.rng.randint(2)
+        points.append([face_x[idx], face_y[idx], z_dim(z)])
+    return [np.array(points).T, dp[1]]
+
+  def __iter__(self):
+    for dp in self.incoming_df:
+      dp = self.map(dp, self.num)
+      yield dp
+
+
+class Model(ModelDesc):
+  def inputs(self):
+    """Inputs are
+    - pointcloud [batch, dim_position, num_points]
+    - label [batch]
+    """
+    return [tf.placeholder(tf.float32, (None, PC['dp'], PC['num']), 'positions'),
+            tf.placeholder(tf.int32, (None,), 'label')]
+
+  def build_graph(self, positions, label):
+
+    positions = positions / 16. - 1
+    # initial features are the position them self
+    features = positions
+    neighbors = knn_bruteforce(positions, K=16)
+
+    x = features
+
+    def subsample(x):
+      # probably too simplistic, just kick out 3 of 4 points randomly
+      # see our paper IDISS approach in the paper for better sub-sampling
+      n = x.shape.as_list()[-1]
+      return x[:, :, :n // 4]
+
+    # similar to traditional networks
+    for stage in range(4):
+      if stage > 0:
+        x = flex_pooling(x, neighbors)
+        x = subsample(x)
+        positions = subsample(positions)
+        neighbors = knn_bruteforce(positions, K=16)
+
+      x = flex_convolution(x, positions, neighbors, 64 *
+                           (stage + 1), activation=tf.nn.relu)
+      x = flex_convolution(x, positions, neighbors, 64 *
+                           (stage + 1), activation=tf.nn.relu)
+
+    if USE_POOLING:
+      # either do max-pooling of all remaining points...
+      x = tf.expand_dims(x, axis=-1)
+      x = tf.layers.max_pooling2d(x, [1, 16], [1, 16])
+    else:
+      # ... or do a flex-conv in (0, 0, 0) with all points as neighbors
+      positions = tf.concat([positions, positions[:, :, :1] * 0], axis=-1)
+      x = tf.concat([x, x[:, :, :1] * 0], axis=-1)
+      K = positions.shape.as_list()[-1]
+      neighbors = knn_bruteforce(positions, K=K)
+      x = flex_convolution(x, positions, neighbors, 1024, activation=tf.nn.relu)
+      x = x[:, :, -1:]
+
+    # from now on just the code part we copied from the Tensorpack framework
+    x = tf.layers.flatten(x)
+    x = tf.layers.dense(x, 512, activation=tf.nn.relu, name='fc0')
+    logits = tf.layers.dense(x, 10, activation=tf.identity, name='fc1')
+
+    cost = tf.nn.sparse_softmax_cross_entropy_with_logits(
+        logits=logits, labels=label)
+    cost = tf.reduce_mean(cost, name='cross_entropy_loss')
+
+    correct = tf.cast(tf.nn.in_top_k(logits, label, 1),
+                      tf.float32, name='correct')
+    accuracy = tf.reduce_mean(correct, name='accuracy')
+
+    train_error = tf.reduce_mean(1 - correct, name='train_error')
+    summary.add_moving_summary(train_error, accuracy)
+    return cost
+
+  def optimizer(self):
+    # nothing fancy here, just stick with the defaults
+    return tf.train.AdamOptimizer(1e-4)
+
+
+def get_data():
+  df_train = dataset.Mnist('train')
+  df_train = Digit2Cloud(df_train, num=PC['num'])
+  df_train = PrefetchDataZMQ(df_train, 2)
+  df_train = BatchData(df_train, BATCH_SIZE)
+
+  df_val = dataset.Mnist('test')
+  df_val = Digit2Cloud(df_val, num=PC['num'])
+  df_val = PrefetchDataZMQ(df_val, 2)
+  df_val = BatchData(df_val, BATCH_SIZE)
+  return df_train, df_val
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--gpu', help='comma separated list of GPU(s) to use.')
+  parser.add_argument('--load', help='load model')
+  parser.add_argument('--fusion', help='run sampling', default='',
+                      choices=['pooling', 'conv'])
+  args = parser.parse_args()
+
+  if args.gpu:
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
+
+  logger.set_logger_dir('train_log/fusion_%s' % (args.fusion))
+
+  dataset_train, dataset_test = get_data()
+  steps_per_epoch = len(dataset_train)
+
+  USE_POOLING = (args.fusion == 'pooling')
+
+  # get the config which contains everything necessary in a training
+  config = TrainConfig(
+      model=Model(),
+      # The input source for training. FeedInput is slow, this is just for demo purpose.
+      # In practice it's best to use QueueInput or others. See tutorials for details.
+      data=FeedInput(dataset_train),
+      callbacks=[
+          ModelSaver(),       # save the model after every epoch
+          InferenceRunner(    # run inference(for validation) after every epoch
+              dataset_test,   # the DataFlow instance used for validation
+              ScalarStats(['cross_entropy_loss', 'accuracy'])),
+      ],
+      extra_callbacks=[
+          MovingAverageSummary(),
+          ProgressBar(['accuracy', 'cross_entropy_loss']),
+          MergeAllSummaries(),
+          RunUpdateOps()
+      ],
+      steps_per_epoch=steps_per_epoch,
+      max_epoch=100,
+  )
+  launch_train_with_config(config, SimpleTrainer())