[SIM/TF2] Add epoch tracking in checkpoints.

Author: tomsiadev

TensorFlow2/Recommendation/SIM/main.py

@@ -38,7 +38,8 @@
 def init_checkpoint_manager(model, optimizer, save_checkpoint_path, load_checkpoint_path):
     checkpoint = tf.train.Checkpoint(
         model=model,
-        optimizer=optimizer
+        optimizer=optimizer,
+        epoch=tf.Variable(-1, name='epoch')
     )
 
     checkpoint_manager = tf.train.CheckpointManager(
@@ -283,7 +284,7 @@ def model_step(batch, model, model_fn, optimizer, amp, first_batch):
     return loss_dict
 
 
-def run_single_epoch(model, model_fn, data_iterator, optimizer, amp, epoch, benchmark, performance_calculator):
+def run_single_epoch(model, model_fn, data_iterator, optimizer, amp, start_epoch, epoch, benchmark, performance_calculator):
 
     for current_step, batch in enumerate(data_iterator):
         if benchmark and performance_calculator.completed:
@@ -296,7 +297,7 @@ def run_single_epoch(model, model_fn, data_iterator, optimizer, amp, epoch, benc
         n_samples = len(batch[1])
         step_throughput = performance_calculator(n_samples)
         step_dict["samples/s"] = step_throughput
-        dllogger.log(data=step_dict, step=(epoch, current_step))
+        dllogger.log(data=step_dict, step=(start_epoch + epoch, current_step))
 
 
 def train(model, model_fn, data_iterator_train, data_iterator_test, optimizer, amp, epochs,
@@ -307,8 +308,10 @@ def train(model, model_fn, data_iterator_train, data_iterator_test, optimizer, a
 
     all_epochs_results = []
 
-    for epoch in range(epochs):
-        run_single_epoch(model, model_fn, data_iterator_train, optimizer, amp, epoch, benchmark, performance_calculator)
+    start_epoch = checkpoint_manager.checkpoint.epoch.numpy().item() + 1
+
+    for epoch in range(epochs - start_epoch):
+        run_single_epoch(model, model_fn, data_iterator_train, optimizer, amp, start_epoch, epoch, benchmark, performance_calculator)
 
         if not benchmark:
             # we dump throughput results for consecutive epochs for a regular training job (w/o --benchmark flag)
@@ -321,10 +324,11 @@ def train(model, model_fn, data_iterator_train, data_iterator_test, optimizer, a
             results_data.update(results_eval_test)
 
             if save_checkpoint:
+                checkpoint_manager.checkpoint.epoch.assign(epoch)
                 checkpoint_manager.save()
 
             if hvd.rank() == 0:
-                dllogger.log(data=results_data, step=(epoch,))
+                dllogger.log(data=results_data, step=(start_epoch + epoch,))
 
             performance_calculator.init()  # restart for another epoch
 

TensorFlow2/Recommendation/SIM/scripts/run_model.sh

@@ -28,6 +28,8 @@ Usage: bash scripts/run_model.sh
 --batch_size            Batch size.
 --results_dir           Path to output directory. Default: /tmp/sim.
 --log_filename          Name of output log file within results_dir. Default: log.json.
+--save_checkpoint_path  Path to output checkpoint after training.
+--load_checkpoint_path  Path from which to restore checkpoint for inference or suspend/resume training.
 EOF
 }
 
@@ -78,9 +80,12 @@ batch_size_option=$(get_option_or_use_default --global_batch_size $batch_size)
 epochs_option=$(get_option_or_use_default --epochs $epochs)
 results_dir_option=$(get_option_or_use_default --results_dir $results_dir)
 log_filename_option=$(get_option_or_use_default --log_filename $log_filename)
+save_checkpoint_path_option=$(get_option_or_use_default --save_checkpoint_path $save_checkpoint_path)
+load_checkpoint_path_option=$(get_option_or_use_default --load_checkpoint_path $load_checkpoint_path)
 
 command="mpiexec --allow-run-as-root --bind-to socket -np ${gpus} python main.py --dataset_dir ${data_path} --drop_remainder ${epochs_option} 
-${xla_arg} ${amp_arg} ${benchmark_arg} ${mode_option} ${benchmark_steps_option} ${batch_size_option} ${results_dir_option} ${log_filename_option}"
+${xla_arg} ${amp_arg} ${benchmark_arg} ${mode_option} ${benchmark_steps_option} ${batch_size_option} ${results_dir_option} ${log_filename_option}
+${save_checkpoint_path_option} ${load_checkpoint_path_option}"
 
 printf "[INFO] Running:\n%s\n" "${command}"
 # run