Unambiguous compression setup to resume properly

examples/tensorflow/classification/main.py

@@ -23,6 +23,8 @@
 from nncf.tensorflow import create_compression_callbacks
 from nncf.tensorflow.helpers.model_manager import TFOriginalModelManager
 from nncf.tensorflow.initialization import register_default_init_args
+from nncf.tensorflow.utils.state import TFCompressionState
+from nncf.tensorflow.utils.state import TFCompressionStateLoader
 
 from examples.tensorflow.classification.datasets.builder import DatasetBuilder
 from examples.tensorflow.common.argparser import get_common_argument_parser
@@ -131,6 +133,12 @@ def resume_from_checkpoint(checkpoint, ckpt_path, steps_per_epoch):
     return initial_epoch
 
 
+def load_compression_state(ckpt_path: str):
+    checkpoint = tf.train.Checkpoint(compression_state=TFCompressionStateLoader())
+    load_checkpoint(checkpoint, ckpt_path)
+    return checkpoint.compression_state.state
+
+
 def run(config):
     strategy = get_distribution_strategy(config)
     if config.metrics_dump is not None:
@@ -168,13 +176,14 @@ def run(config):
                 return_dict=True)
             uncompressed_model_accuracy = 100 * results['acc@1']
 
+    compression_state = None
+    if resume_training:
+        compression_state = load_compression_state(config.ckpt_path)
+
     with TFOriginalModelManager(model_fn, **model_params) as model:
         with strategy.scope():
-            compression_ctrl, compress_model = create_compressed_model(model,
-                                                                       nncf_config,
-                                                                       should_init=not resume_training)
-            compression_callbacks = create_compression_callbacks(compression_ctrl,
-                                                                 log_dir=config.log_dir)
+            compression_ctrl, compress_model = create_compressed_model(model, nncf_config, compression_state)
+            compression_callbacks = create_compression_callbacks(compression_ctrl, log_dir=config.log_dir)
 
             scheduler = build_scheduler(
                 config=config,
@@ -201,7 +210,8 @@ def run(config):
 
             compress_model.summary()
 
-            checkpoint = tf.train.Checkpoint(model=compress_model, compression_ctrl=compression_ctrl)
+            checkpoint = tf.train.Checkpoint(model=compress_model,
+                                             compression_state=TFCompressionState(compression_ctrl))
 
             initial_epoch = 0
             if resume_training:
@@ -279,9 +289,12 @@ def export(config):
                                     pretrained=config.get('pretrained', False),
                                     weights=config.get('weights', None))
     model = model(**model_params)
-    compression_ctrl, compress_model = create_compressed_model(model,
-                                                               config.nncf_config,
-                                                               should_init=False)
+
+    compression_state = None
+    if config.ckpt_path:
+        compression_state = load_compression_state(config.ckpt_path)
+
+    compression_ctrl, compress_model = create_compressed_model(model, config.nncf_config, compression_state)
 
     metrics = [
         tf.keras.metrics.CategoricalAccuracy(name='acc@1'),
@@ -293,7 +306,8 @@ def export(config):
                            metrics=metrics)
     compress_model.summary()
 
-    checkpoint = tf.train.Checkpoint(model=compress_model, compression_ctrl=compression_ctrl)
+    checkpoint = tf.train.Checkpoint(model=compress_model,
+                                     compression_state=TFCompressionState(compression_ctrl))
 
     if config.ckpt_path is not None:
         load_checkpoint(checkpoint=checkpoint,

examples/tensorflow/common/prepare_checkpoint.py

@@ -16,6 +16,8 @@
 import tensorflow as tf
 
 from nncf.tensorflow.helpers.model_creation import create_compressed_model
+from nncf.tensorflow.utils.state import TFCompressionState
+from nncf.tensorflow.utils.state import TFCompressionStateLoader
 from examples.tensorflow.common.logger import logger
 from examples.tensorflow.common.sample_config import create_sample_config
 from examples.tensorflow.common.argparser import get_common_argument_parser
@@ -71,18 +73,24 @@ def load_checkpoint(checkpoint, ckpt_path):
     return None
 
 
+def load_compression_state(ckpt_path: str):
+    checkpoint = tf.train.Checkpoint(compression_state=TFCompressionStateLoader())
+    load_checkpoint(checkpoint, ckpt_path)
+    return checkpoint.compression_state.state
+
+
 def od_checkpoint_saver(config):
     """
     Load object detection checkpoint and re-save it without optimizer (memory footprint is reduced).
     """
     model_builder = get_model_od_builder(config)
     model = model_builder.build_model()
 
-    compression_ctrl, compress_model = create_compressed_model(model,
-                                                               config.nncf_config,
-                                                               should_init=False)
+    compression_state = load_compression_state(config.ckpt_path)
+    compression_ctrl, compress_model = create_compressed_model(model, config.nncf_config, compression_state)
 
-    checkpoint = tf.train.Checkpoint(model=compress_model, compression_ctrl=compression_ctrl)
+    checkpoint = tf.train.Checkpoint(model=compress_model,
+                                     compression_state=TFCompressionState(compression_ctrl))
     load_and_save_checkpoint(checkpoint, config)
 
 
@@ -93,13 +101,12 @@ def seg_checkpoint_saver(config):
     model_builder = get_model_seg_builder(config)
     model = model_builder.build_model()
 
-    compression_ctrl, compress_model = create_compressed_model(model,
-                                                               config.nncf_config,
-                                                               should_init=False)
+    compression_state = load_compression_state(config.ckpt_path)
+    compression_ctrl, compress_model = create_compressed_model(model, config.nncf_config, compression_state)
 
     variables = get_variables(compress_model)
     checkpoint = tf.train.Checkpoint(variables=variables,
-                                     compression_ctrl=compression_ctrl,
+                                     compression_state=TFCompressionState(compression_ctrl),
                                      step=tf.Variable(0))
     load_and_save_checkpoint(checkpoint, config)
 

examples/tensorflow/object_detection/main.py

@@ -25,6 +25,8 @@
 from nncf.common.utils.tensorboard import prepare_for_tensorboard
 from nncf.config.utils import is_accuracy_aware_training
 from nncf.config.structures import ModelEvaluationArgs
+from nncf.tensorflow.utils.state import TFCompressionState
+from nncf.tensorflow.utils.state import TFCompressionStateLoader
 
 from examples.tensorflow.common.argparser import get_common_argument_parser
 from examples.tensorflow.common.distributed import get_distribution_strategy
@@ -120,6 +122,12 @@ def resume_from_checkpoint(checkpoint_manager, ckpt_path, steps_per_epoch):
     return initial_epoch, initial_step
 
 
+def load_compression_state(ckpt_path: str):
+    checkpoint = tf.train.Checkpoint(compression_state=TFCompressionStateLoader())
+    load_checkpoint(checkpoint, ckpt_path)
+    return checkpoint.compression_state.state
+
+
 def create_test_step_fn(strategy, model, predict_post_process_fn):
     """Creates a distributed test step"""
 
@@ -271,10 +279,9 @@ def run(config):
         write_metrics(0, config.metrics_dump)
 
     # Create dataset
-    builders = get_dataset_builders(config, strategy.num_replicas_in_sync)
-    datasets = [builder.build() for builder in builders]
-    train_builder, test_builder = builders
-    train_dataset, test_dataset = datasets
+    train_builder, test_builder = get_dataset_builders(config, strategy.num_replicas_in_sync)
+    train_dataset = train_builder.build()
+    test_dataset = test_builder.build()
     train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)
     test_dist_dataset = strategy.experimental_distribute_dataset(test_dataset)
 
@@ -300,13 +307,15 @@ def model_eval_fn(model):
 
     resume_training = config.ckpt_path is not None
 
+    compression_state = None
+    if resume_training:
+        compression_state = load_compression_state(config.ckpt_path)
+
     with TFOriginalModelManager(model_builder.build_model,
                                 weights=config.get('weights', None)) as model:
         with strategy.scope():
             config.nncf_config.register_extra_structs([ModelEvaluationArgs(eval_fn=model_eval_fn)])
-            compression_ctrl, compress_model = create_compressed_model(model,
-                                                                       nncf_config,
-                                                                       should_init=not resume_training)
+            compression_ctrl, compress_model = create_compressed_model(model, nncf_config, compression_state)
             scheduler = build_scheduler(
                 config=config,
                 steps_per_epoch=steps_per_epoch)
@@ -321,7 +330,7 @@ def model_eval_fn(model):
 
             checkpoint = tf.train.Checkpoint(model=compress_model,
                                              optimizer=optimizer,
-                                             compression_ctrl=compression_ctrl)
+                                             compression_state=TFCompressionState(compression_ctrl))
             checkpoint_manager = tf.train.CheckpointManager(checkpoint, config.checkpoint_save_dir, max_to_keep=None)
 
             initial_epoch = initial_step = 0
@@ -358,7 +367,6 @@ def validate_fn(model, epoch):
                                                          validate_fn=validate_fn,
                                                          tensorboard_writer=config.tb,
                                                          log_dir=config.log_dir)
-
         else:
             train(train_step, test_step, eval_metric, train_dist_dataset, test_dist_dataset,
                   initial_epoch, initial_step, epochs, steps_per_epoch, checkpoint_manager,
@@ -382,12 +390,15 @@ def export(config):
     model_builder = get_model_builder(config)
     model = model_builder.build_model(weights=config.get('weights', None))
 
-    compression_ctrl, compress_model = create_compressed_model(model,
-                                                               config.nncf_config,
-                                                               should_init=False)
+    compression_state = None
+    if config.ckpt_path:
+        compression_state = load_compression_state(config.ckpt_path)
+
+    compression_ctrl, compress_model = create_compressed_model(model, config.nncf_config, compression_state)
 
     if config.ckpt_path:
-        checkpoint = tf.train.Checkpoint(model=compress_model)
+        checkpoint = tf.train.Checkpoint(model=compress_model,
+                                         compression_state=TFCompressionState(compression_ctrl))
         load_checkpoint(checkpoint, config.ckpt_path)
 
     save_path, save_format = get_saving_parameters(config)

examples/tensorflow/segmentation/evaluation.py

@@ -16,7 +16,10 @@
 import tensorflow as tf
 
 from nncf.tensorflow import create_compressed_model
+from nncf.tensorflow import register_default_init_args
 from nncf.tensorflow.helpers.model_manager import TFOriginalModelManager
+from nncf.tensorflow.utils.state import TFCompressionState
+from nncf.tensorflow.utils.state import TFCompressionStateLoader
 
 from examples.tensorflow.common.argparser import get_common_argument_parser
 from examples.tensorflow.common.distributed import get_distribution_strategy
@@ -90,8 +93,12 @@ def get_dataset_builders(config, num_devices):
     val_builder = COCODatasetBuilder(config=config,
                                      is_train=False,
                                      num_devices=num_devices)
-
-    return val_builder
+    config_ = config.deepcopy()
+    config_.batch_size = val_builder.batch_size
+    calibration_builder = COCODatasetBuilder(config=config_,
+                                             is_train=True,
+                                             num_devices=1)
+    return val_builder, calibration_builder
 
 
 def load_checkpoint(checkpoint, ckpt_path):
@@ -115,6 +122,12 @@ def load_checkpoint(checkpoint, ckpt_path):
     return None
 
 
+def load_compression_state(ckpt_path: str):
+    checkpoint = tf.train.Checkpoint(compression_state=TFCompressionStateLoader())
+    load_checkpoint(checkpoint, ckpt_path)
+    return checkpoint.compression_state.state
+
+
 def evaluate(test_step, metric, test_dist_dataset, num_batches, print_freq):
     """Runs evaluation steps and aggregate metrics"""
     timer = Timer()
@@ -164,41 +177,56 @@ def test_step(dataset_inputs):
     return test_step
 
 
-def run_evaluation(config, eval_timeout=None):
-    """Runs evaluation on checkpoint save directory"""
-    strategy = get_distribution_strategy(config)
-    if config.metrics_dump is not None:
-        write_metrics(0, config.metrics_dump)
-
-    dataset_builder = get_dataset_builders(config, strategy.num_replicas_in_sync)
-    dataset = dataset_builder.build()
-    num_batches = dataset_builder.steps_per_epoch
-    test_dist_dataset = strategy.experimental_distribute_dataset(dataset)
-
-    # We use `model_batch_size` to create input layer for model
-    config.model_batch_size = dataset_builder.batch_size
-
-    model_builder = get_model_builder(config)
+def restore_compressed_model(config, strategy, model_builder, ckpt_path = None):
+    compression_state = None
+    if ckpt_path:
+        compression_state = load_compression_state(ckpt_path)
 
     with TFOriginalModelManager(model_builder.build_model,
                                 weights=config.get('weights', None),
                                 is_training=False) as model:
         with strategy.scope():
             compression_ctrl, compress_model = create_compressed_model(model,
                                                                        config.nncf_config,
-                                                                       should_init=False)
+                                                                       compression_state)
+
             variables = get_variables(compress_model)
             checkpoint = tf.train.Checkpoint(variables=variables,
-                                             compression_ctrl=compression_ctrl,
+                                             compression_state=TFCompressionState(compression_ctrl),
                                              step=tf.Variable(0))
-            eval_metric = model_builder.eval_metrics()
-            predict_post_process_fn = model_builder.post_processing
+            if ckpt_path:
+                load_checkpoint(checkpoint, config.ckpt_path)
 
-    test_step = create_test_step_fn(strategy, compress_model, predict_post_process_fn)
+    return compression_ctrl, compress_model, checkpoint
+
+
+def run_evaluation(config, eval_timeout=None):
+    """Runs evaluation on checkpoint save directory"""
+    strategy = get_distribution_strategy(config)
+    if config.metrics_dump is not None:
+        write_metrics(0, config.metrics_dump)
+
+    validation_builder, calibration_builder = get_dataset_builders(config, strategy.num_replicas_in_sync)
+    calibration_dataset = calibration_builder.build()
+    val_dataset = validation_builder.build()
+    num_batches = validation_builder.steps_per_epoch
+    test_dist_dataset = strategy.experimental_distribute_dataset(val_dataset)
+
+    config.nncf_config = register_default_init_args(nncf_config=config.nncf_config,
+                                                    data_loader=calibration_dataset,
+                                                    batch_size=validation_builder.global_batch_size)
+
+    # We use `model_batch_size` to create input layer for model
+    config.model_batch_size = validation_builder.batch_size
+
+    model_builder = get_model_builder(config)
+    eval_metric = model_builder.eval_metrics()
+    predict_post_process_fn = model_builder.post_processing
 
     if 'test' in config.mode:
-        if config.ckpt_path:
-            load_checkpoint(checkpoint, config.ckpt_path)
+        compression_ctrl, compress_model, _ = restore_compressed_model(config, strategy, model_builder,
+                                                                       config.ckpt_path)
+        test_step = create_test_step_fn(strategy, compress_model, predict_post_process_fn)
 
         statistics = compression_ctrl.statistics()
         logger.info(statistics.to_str())
@@ -213,12 +241,17 @@ def run_evaluation(config, eval_timeout=None):
 
     elif 'train' in config.mode:
         validation_summary_writer = SummaryWriter(config.log_dir, 'validation')
-        checkpoint_dir = config.checkpoint_save_dir
-        eval_timeout = config.eval_timeout
 
-        for checkpoint_path in tf.train.checkpoints_iterator(checkpoint_dir, timeout=eval_timeout):
-            status = checkpoint.restore(checkpoint_path)
-            status.expect_partial()
+        is_first_checkpoint = True
+        for checkpoint_path in tf.train.checkpoints_iterator(config.checkpoint_save_dir, config.eval_timeout):
+            if is_first_checkpoint:
+                is_first_checkpoint = False
+                _, compress_model, checkpoint = restore_compressed_model(config, strategy, model_builder,
+                                                                         checkpoint_path)
+                test_step = create_test_step_fn(strategy, compress_model, predict_post_process_fn)
+            else:
+                checkpoint.restore(checkpoint_path).expect_partial()
+
             logger.info('Checkpoint file {} found and restoring from checkpoint'.format(checkpoint_path))
             logger.info('Checkpoint step: {}'.format(checkpoint.step.numpy()))
             metric_result = evaluate(test_step, eval_metric, test_dist_dataset, num_batches, config.print_freq)
@@ -238,17 +271,8 @@ def run_evaluation(config, eval_timeout=None):
 def export(config):
     model_builder = get_model_builder(config)
 
-    with TFOriginalModelManager(model_builder.build_model,
-                                weights=config.get('weights', None),
-                                is_training=False) as model:
-        compression_ctrl, compress_model = create_compressed_model(model,
-                                                                   config.nncf_config,
-                                                                   should_init=False)
-
-    if config.ckpt_path:
-        variables = get_variables(compress_model)
-        checkpoint = tf.train.Checkpoint(variables=variables)
-        load_checkpoint(checkpoint, config.ckpt_path)
+    strategy = tf.distribute.get_strategy()
+    compression_ctrl, _, _ = restore_compressed_model(config, strategy, model_builder, config.ckpt_path)
 
     save_path, save_format = get_saving_parameters(config)
     compression_ctrl.export_model(save_path, save_format)