Workflow

API Call Sequence

The following figure shows the API call sequence. The training runs on the CPU or GPU environment of the TensorFlow framework. Based on the inference script of the open-source framework, the AMCT API is called to compress the model. The compressed model must be converted into an offline model that adapts to the AI processor using the ATC before it can be used for inference on the AI processor.

Figure 1 API call sequence for compression combination

The user implements the operations in blue, while those in gray are implemented by using AMCT APIs. Specifically, import the package to the source TensorFlow network inference code and call APIs where appropriate for compression.

Example

1. Take the following steps to get started. Update the sample code based on your situation.
2. Tweak the arguments passed to AMCT API calls as required. QAT relies on the user training result. Ensure that a TensorFlow training script that yields satisfactory training accuracy is available.

1. Import the AMCT package and set the log level.

   1 2	import amct_tensorflow as amct amct.set_logging_level(print_level='info', save_level='info')

2. (Optional) Build a graph, read the trained parameters, and run inference on the graph in the TensorFlow environment to validate the inference script and environment setup. (Update the sample code based on your situation.)

   This step is recommended as it guarantees a properly functioning original model for inference with acceptable accuracy. You can use a subset from the test dataset to improve the efficiency.

   1	user_test_evaluate_model(evaluate_model, test_data)

3. Build a training graph. (Update the sample code based on your situation.)

   1	train_graph = user_load_train_graph()

4. Run AMCT to perform training with the sparsity operator and quantization parameters.
   1. Modify the graph by inserting sparsity and quantization operators into the graph.

      Before compression, call the create_compressed_retrain_model API to modify the trained graph based on the simplified configuration file and the original model. Specifically, insert the filter-level sparsity (or 2:4 structured sparsity) and QAT operators to generate the graph for compression combination (corresponding to 1 in Figure 1).

      1 2 3 4 5

      record_file = './tmp/record.txt' retrain_ops = amct.create_compressed_retrain_model(graph=train_graph, config_defination=simple_cfg, outputs=user_model_outputs, record_file=record_file)

   2. Implement gradient descent optimization on the modified graph, train the graph on the training dataset, and train model parameters and quantization factors. (Update the sample code based on your situation.)
      1. Implement gradient descent optimization on the modified graph. Perform this step after 4.a.
         Call RMSPropOptimizer to implement gradient descent optimization.

         1 2 3

         optimizer = tf.compat.v1.train.RMSPropOptimizer( ARGS.learning_rate, momentum=ARGS.momentum) train_op = optimizer.minimize(loss)

      2. Create a session to train the model, and save the trained parameters as a checkpoint file (corresponding to 2 and 3 in Figure 1).
         Note: Restore the model parameters from existing checkpoints and then train the model. The parameters saved during training should include quantization factors. Quantization factors are generated after the first batch_num training. If the number of training times is less than batch_num, the training fails.

         1 2 3 4 5

         with tf.Session() as sess: sess.run(tf.compat.v1.global_variables_initializer()) sess.run(outputs) # Save the trained parameters as a checkpoint file. saver_save.save(sess, retrain_ckpt, global_step=0)

5. Build an inference graph. (Update the sample code based on your situation.)

   1	test_graph = user_load_test_graph()

6. Run AMCT to perform compression combination.
   1. Modify the inference graph.

      Construct an inference graph (that is, set is_training to False for the BN), and then call the create_compressed_retrain_model API to modify the graph before compression combination based on the quantization configuration file, that is, to insert filter-level sparsity (or 2:4 structured sparsity) and QAT operators into the graph for subsequent model freezing and inference, and generate a sparsity and quantization factor record file (corresponding to 4 and 5 in Figure 1).

   2. Create a session to restore the training parameters, infer the output node (retrain_ops[-1]) of the quantized model, write the quantization factors to the record file, and freeze the inference graph into a .pb model (corresponding to 6 and 7 in Figure 1). (Update the sample code based on your situation.)
      Note: The parameters to be inferred (the output tensor of retrain_ops[-1]) and to be restored must be in the same session. The .pb model generated from an inference graph contains the trained parameter values.

      1 2 3 4 5 6 7 8 9 10 11 12 13

      variables_to_restore = tf.compat.v1.global_variables() saver_restore = tf.compat.v1.train.Saver(variables_to_restore) with tf.Session() as sess: sess.run(tf.compat.v1.global_variables_initializer()) # Restore training parameters. saver_restore.restore(sess, retrain_ckpt) # Write quantization factors to the record file. Note that you can skip this step if quantization is not involved. sess.run(retrain_ops[-1]) # Save the model as a .pb model. constant_graph = tf.compat.v1.graph_util.convert_variables_to_constants( sess, eval_graph.as_graph_def(), [output.name[:-2] for output in outputs]) with tf.io.gfile.GFile(frozen_quant_eval_pb, 'wb') as f: f.write(constant_graph.SerializeToString())

   3. Save the compressed model.
      Remove the sparsity operator, insert quantization operators such as AscendQuant and AscendDequant into the frozen model, and save the compressed model based on the sparsity and quantization factor record file.

      1 2 3 4 5

      compressed_model_path = './result/user_model' amct.save_compressed_retrain_model(pb_model=trained_pb, outputs=user_model_outputs, record_file=record_file, save_path=compressed_model_path)

7. (Optional) Run inference on the compressed model user_model_compressed.pb in the TensorFlow environment based on the test dataset to test the accuracy. (Update the sample code based on your situation.)
   Compare the accuracy of the fake-quantized model with that of the original model (see 2).

   1 2	compressed_model = './results/user_model_compressed.pb' user_do_inference(compressed_model, test_data)