starting epoch

tensorpack/callbacks/base.py

@@ -29,7 +29,7 @@ class Callback(object):
     def before_train(self, trainer):
         self.trainer = trainer
         self.graph = tf.get_default_graph()
-        self.epoch_num = 0
+        self.epoch_num = self.trainer.config.starting_epoch
         self._before_train()
 
     def _before_train(self):
@@ -59,8 +59,8 @@ class Callback(object):
         """
         epoch_num is the number of epoch finished.
         """
-        self.epoch_num += 1
         self._trigger_epoch()
+        self.epoch_num += 1
 
     def _trigger_epoch(self):
         """

tensorpack/train/base.py

@@ -77,7 +77,7 @@ class Trainer(object):
                 self.global_step = get_global_step()
                 logger.info("Start training with global_step={}".format(self.global_step))
 
-                for epoch in range(1, self.config.max_epoch):
+                for epoch in range(self.config.starting_epoch, self.config.max_epoch):
                     with timed_operation(
                         'Epoch {}, global_step={}'.format(
                             epoch, self.global_step + self.config.step_per_epoch)):

tensorpack/train/config.py

@@ -29,8 +29,8 @@ class TrainConfig(object):
             session_init: a tensorpack.utils.sessinit.SessionInit instance to
                 initialize variables of a session. default to a new session.
             model: a ModelDesc instance
-            step_per_epoch: the number of steps (parameter updates) to perform
-                in each epoch.
+            starting_epoch: int. default to be 1.
+            step_per_epoch: the number of steps (SGD updates) to perform in each epoch.
             max_epoch: maximum number of epoch to run training. default to 100
             nr_tower: int. number of towers. default to 1.
         """
@@ -50,6 +50,7 @@ class TrainConfig(object):
         self.session_init = kwargs.pop('session_init', NewSession())
         assert_type(self.session_init, SessionInit)
         self.step_per_epoch = int(kwargs.pop('step_per_epoch'))
+        self.starting_epoch = int(kwargs.pop('starting_epoch', 1))
         self.max_epoch = int(kwargs.pop('max_epoch', 100))
         assert self.step_per_epoch > 0 and self.max_epoch > 0
         self.nr_tower = int(kwargs.pop('nr_tower', 1))