Use Chain to refactor trainer.

reyoung · reyoung · commit aba84aad8370 · 2016-12-23T14:13:54.000+08:00
diff --git a/demo/mnist/api_train.py b/demo/mnist/api_train.py
@@ -14,6 +14,8 @@
 from mnist_util import read_from_mnist
 from paddle.trainer_config_helpers import *
 
+from trainer import *
+
 
 def optimizer_config():
     settings(
@@ -72,122 +74,132 @@ def input_order_converter(generator):
         yield each_item['pixel'], each_item['label']
 
 
-def main():
-    api.initPaddle("-use_gpu=false", "-trainer_count=4")  # use 4 cpu cores
-
-    # get enable_types for each optimizer.
-    # enable_types = [value, gradient, momentum, etc]
-    # For each optimizer(SGD, Adam), GradientMachine should enable different
-    # buffers.
-    opt_config_proto = parse_optimizer_config(optimizer_config)
-    opt_config = api.OptimizationConfig.createFromProto(opt_config_proto)
-    _temp_optimizer_ = api.ParameterOptimizer.create(opt_config)
-    enable_types = _temp_optimizer_.getParameterTypes()
-
-    # Create Simple Gradient Machine.
-    model_config = parse_network_config(network_config)
-    m = api.GradientMachine.createFromConfigProto(
-        model_config, api.CREATE_MODE_NORMAL, enable_types)
-
-    # This type check is not useful. Only enable type hint in IDE.
-    # Such as PyCharm
-    assert isinstance(m, api.GradientMachine)
-
-    # Initialize Parameter by numpy.
-    init_parameter(network=m)
-
-    # Create Local Updater. Local means not run in cluster.
-    # For a cluster training, here we can change to createRemoteUpdater
-    # in future.
-    updater = api.ParameterUpdater.createLocalUpdater(opt_config)
-    assert isinstance(updater, api.ParameterUpdater)
-
-    # Initialize ParameterUpdater.
-    updater.init(m)
-
-    # DataProvider Converter is a utility convert Python Object to Paddle C++
-    # Input. The input format is as same as Paddle's DataProvider.
-    converter = DataProviderConverter(
-        input_types=[dp.dense_vector(784), dp.integer_value(10)])
-
-    train_file = './data/raw_data/train'
-    test_file = './data/raw_data/t10k'
-
-    # start gradient machine.
-    # the gradient machine must be started before invoke forward/backward.
-    # not just for training, but also for inference.
-    m.start()
-
-    # evaluator can print error rate, etc. It is a C++ class.
-    batch_evaluator = m.makeEvaluator()
-    test_evaluator = m.makeEvaluator()
-
-    # Get Train Data.
-    # TrainData will stored in a data pool. Currently implementation is not care
-    # about memory, speed. Just a very naive implementation.
-    train_data_generator = input_order_converter(read_from_mnist(train_file))
-    train_data = BatchPool(train_data_generator, 512)
-
-    # outArgs is Neural Network forward result. Here is not useful, just passed
-    # to gradient_machine.forward
-    outArgs = api.Arguments.createArguments(0)
-
-    for pass_id in xrange(2):  # we train 2 passes.
-        updater.startPass()
-
-        for batch_id, data_batch in enumerate(train_data()):
-            # data_batch is input images.
-            # here, for online learning, we could get data_batch from network.
-
-            # Start update one batch.
-            pass_type = updater.startBatch(len(data_batch))
-
-            # Start BatchEvaluator.
-            # batch_evaluator can be used between start/finish.
-            batch_evaluator.start()
-
-            # forwardBackward is a shortcut for forward and backward.
-            # It is sometimes faster than invoke forward/backward separately,
-            # because in GradientMachine, it may be async.
-            m.forwardBackward(converter(data_batch), outArgs, pass_type)
-
-            for each_param in m.getParameters():
-                updater.update(each_param)
-
-            # Get cost. We use numpy to calculate total cost for this batch.
-            cost_vec = outArgs.getSlotValue(0)
-            cost_vec = cost_vec.copyToNumpyMat()
-            cost = cost_vec.sum() / len(data_batch)
-
-            # Make evaluator works.
-            m.eval(batch_evaluator)
-
-            # Print logs.
-            print 'Pass id', pass_id, 'Batch id', batch_id, 'with cost=', \
-                cost, batch_evaluator
-
-            batch_evaluator.finish()
-            # Finish batch.
-            #  * will clear gradient.
-            #  * ensure all values should be updated.
-            updater.finishBatch(cost)
-
+class MonolithicChainItem(RunnerChainItem):
+    def finalize(self, context, next_callback):
+        context.gradient_machine.finish()
+
+    def initialize(self, context, next_callback):
+        api.initPaddle("-use_gpu=false", "-trainer_count=4")  # use 4 cpu cores
+
+        # get enable_types for each optimizer.
+        # enable_types = [value, gradient, momentum, etc]
+        # For each optimizer(SGD, Adam), GradientMachine should enable different
+        # buffers.
+        opt_config_proto = parse_optimizer_config(optimizer_config)
+        opt_config = api.OptimizationConfig.createFromProto(opt_config_proto)
+        _temp_optimizer_ = api.ParameterOptimizer.create(opt_config)
+        enable_types = _temp_optimizer_.getParameterTypes()
+
+        # Create Simple Gradient Machine.
+        model_config = parse_network_config(network_config)
+        context.gradient_machine = api.GradientMachine.createFromConfigProto(
+            model_config, api.CREATE_MODE_NORMAL, enable_types)
+
+        # This type check is not useful. Only enable type hint in IDE.
+        # Such as PyCharm
+        assert isinstance(context.gradient_machine, api.GradientMachine)
+
+        # Initialize Parameter by numpy.
+        init_parameter(network=context.gradient_machine)
+
+        # Create Local Updater. Local means not run in cluster.
+        # For a cluster training, here we can change to createRemoteUpdater
+        # in future.
+        context.updater = api.ParameterUpdater.createLocalUpdater(opt_config)
+        assert isinstance(context.updater, api.ParameterUpdater)
+        context.updater.init(context.gradient_machine)
+
+        # DataProvider Converter is a utility convert Python Object to Paddle C++
+        # Input. The input format is as same as Paddle's DataProvider.
+        context.data_converter = DataProviderConverter(
+            input_types=[dp.dense_vector(784), dp.integer_value(10)])
+
+        train_file = './data/raw_data/train'
+        test_file = './data/raw_data/t10k'
+
+        context.gradient_machine.start()
+
+        # Get Train Data.
+        # TrainData will stored in a data pool. Currently implementation is not care
+        # about memory, speed. Just a very naive implementation.
+        train_data_generator = input_order_converter(
+            read_from_mnist(train_file))
+        train_data = BatchPool(train_data_generator, 512)
+        context.train_data_callback = train_data
+        context.test_file = test_file
+
+        next_callback(context)
+
+    def on_batch_begin(self, context, next_callback):
+        batch_evaluator = context.gradient_machine.makeEvaluator()
+        # outArgs is Neural Network forward result. Here is not useful, just passed
+        # to gradient_machine.forward
+        outArgs = api.Arguments.createArguments(0)
+
+        try:
+            data_batch = next(context.train_data)
+        except StopIteration:
+            return True
+
+        # data_batch is input images.
+        # here, for online learning, we could get data_batch from network.
+
+        # Start update one batch.
+        pass_type = context.updater.startBatch(len(data_batch))
+
+        # Start BatchEvaluator.
+        # batch_evaluator can be used between start/finish.
+        batch_evaluator.start()
+
+        # forwardBackward is a shortcut for forward and backward.
+        # It is sometimes faster than invoke forward/backward separately,
+        # because in GradientMachine, it may be async.
+        context.gradient_machine.forwardBackward(
+            context.data_converter(data_batch), outArgs, pass_type)
+
+        for each_param in context.gradient_machine.getParameters():
+            context.updater.update(each_param)
+
+        # Get cost. We use numpy to calculate total cost for this batch.
+        cost_vec = outArgs.getSlotValue(0)
+        cost_vec = cost_vec.copyToNumpyMat()
+        cost = cost_vec.sum() / len(data_batch)
+
+        # Make evaluator works.
+        context.gradient_machine.eval(batch_evaluator)
+
+        # Print logs.
+        print 'batch with cost=', cost, batch_evaluator
+
+        batch_evaluator.finish()
+        context.cost = cost
+        return False
+
+    def on_pass_begin(self, context, next_callback):
+        context.updater.startPass()
+        context.train_data = context.train_data_callback()
+
+    def on_pass_end(self, context, next_callback):
         # testing stage. use test data set to test current network.
-        updater.apply()
+        outArgs = api.Arguments.createArguments(0)
+        context.updater.apply()
+        test_evaluator = context.gradient_machine.makeEvaluator()
         test_evaluator.start()
-        test_data_generator = input_order_converter(read_from_mnist(test_file))
+        test_data_generator = input_order_converter(
+            read_from_mnist(context.test_file))
         for data_batch in generator_to_batch(test_data_generator, 512):
             # in testing stage, only forward is needed.
-            m.forward(converter(data_batch), outArgs, api.PASS_TEST)
-            m.eval(test_evaluator)
+            context.gradient_machine.forward(
+                context.data_converter(data_batch), outArgs, api.PASS_TEST)
+            context.gradient_machine.eval(test_evaluator)
 
         # print error rate for test data set
-        print 'Pass', pass_id, ' test evaluator: ', test_evaluator
+        print 'Test evaluator: ', test_evaluator
         test_evaluator.finish()
-        updater.restore()
+        context.updater.restore()
 
-        updater.catchUpWith()
-        params = m.getParameters()
+        context.updater.catchUpWith()
+        params = context.gradient_machine.getParameters()
         for each_param in params:
             assert isinstance(each_param, api.Parameter)
             value = each_param.getBuf(api.PARAMETER_VALUE)
@@ -196,9 +208,25 @@ def main():
             # Here, we could save parameter to every where you want
             print each_param.getName(), value
 
-        updater.finishPass()
+        context.updater.finishPass()
+
+    def on_batch_end(self, context, next_callback):
+        # Finish batch.
+        #  * will clear gradient.
+        #  * ensure all values should be updated.
+        context.updater.finishBatch(context.cost)
+        return False
 
-    m.finish()
+    def __init__(self):
+        RunnerChainItem.__init__(self)
+
+
+def main():
+    runner = Runner()
+    runner.add_chain_item(MonolithicChainItem())
+    with runner.use():
+        for _ in xrange(2):
+            runner.run_one_pass()
 
 
 if __name__ == '__main__':
diff --git a/demo/mnist/trainer.py b/demo/mnist/trainer.py
@@ -0,0 +1,109 @@
+import functools
+
+__all__ = ['RunnerChainItem', 'Runner']
+
+
+class RunnerChainItem(object):
+    def __init__(self):
+        pass
+
+    def initialize(self, context, next_callback):
+        next_callback(context)
+
+    def finalize(self, context, next_callback):
+        next_callback(context)
+
+    def on_pass_begin(self, context, next_callback):
+        next_callback(context)
+
+    def on_pass_end(self, context, next_callback):
+        next_callback(context)
+
+    def on_batch_begin(self, context, next_callback):
+        return next_callback(context)
+
+    def on_batch_end(self, context, next_callback):
+        return next_callback(context)
+
+
+def default_next_callback(*args, **kwargs):
+    return False
+
+
+class RunnerContext(object):
+    pass
+
+
+class RunnerSection(object):
+    def __init__(self, runner):
+        self.runner = runner
+
+    def __enter__(self):
+        self.runner.initialize()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.runner.finalize()
+
+
+class Runner(object):
+    def __init__(self):
+        self.chains = []
+
+        self.begin_pass = None
+        self.end_pass = None
+        self.begin_batch = None
+        self.end_batch = None
+        self.finalize = None
+
+        self.context = RunnerContext()
+        self.context.runner = self
+
+    def add_chain_item(self, item):
+        assert isinstance(item, RunnerChainItem)
+        self.chains.append(item)
+
+    def initialize(self):
+        if None not in [
+                self.begin_pass, self.end_pass, self.begin_batch,
+                self.end_batch, self.finalize
+        ]:
+            return False
+        else:
+            assert len(self.chains) != 0
+            actual_init = default_next_callback
+            self.begin_pass = default_next_callback
+            self.end_pass = default_next_callback
+            self.begin_batch = default_next_callback
+            self.end_batch = default_next_callback
+            self.finalize = default_next_callback
+
+            for chain in reversed(self.chains):
+                assert isinstance(chain, RunnerChainItem)
+                actual_init = functools.partial(
+                    chain.initialize, next_callback=actual_init)
+                self.begin_pass = functools.partial(
+                    chain.on_pass_begin, next_callback=self.begin_pass)
+                self.end_pass = functools.partial(
+                    chain.on_pass_end, next_callback=self.end_pass)
+                self.begin_batch = functools.partial(
+                    chain.on_batch_begin, next_callback=self.begin_batch)
+                self.end_batch = functools.partial(
+                    chain.on_batch_end, next_callback=self.end_batch)
+                self.finalize = functools.partial(
+                    chain.finalize, next_callback=self.finalize)
+
+            actual_init(self.context)
+            return True
+
+    def run_one_pass(self):
+        self.begin_pass(self.context)
+        exit_flag = False
+        while not exit_flag:
+            exit_flag = self.begin_batch(self.context)
+            if exit_flag:
+                break
+            exit_flag = self.end_batch(self.context)
+        self.end_pass(self.context)
+
+    def use(self):
+        return RunnerSection(self)
