Tensorflow Post Training Quantization

import tensorflow as tf
import numpy as np
from tensorflow.keras import layers, losses, metrics, optimizers, models

X = tf.random.uniform([100, 1], minval=1, maxval=10.0)
Y = tf.pow(X, 2)

tf.keras.backend.clear_session()

model = models.Sequential()
model.add(layers.Dense(100, input_shape=(1,), activation="relu"))
model.add(layers.Dense(100, activation="relu"))
model.add(layers.Dense(100, activation="relu"))
model.add(layers.Dense(100, activation="relu"))
model.add(layers.Dense(100, activation="relu"))
model.add(layers.Dense(1))

model.summary()

model.compile(optimizer="adam", loss="mse")
history = model.fit(X, Y, batch_size=20, epochs=2000, verbose=0)
print(model.predict([2.0, 10.0, 20]))

model.save("/tmp/pow")

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model("/tmp/pow")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]

tflite = converter.convert()
with open ("/tmp/pow-q.tflite","wb") as f:
    f.write(tflite)
    print("size of tflite-q:", len(tflite))

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node)
  const TfLiteTensor* filter = GetInput(context, node, kWeightsTensor);
  switch (filter->type)
    case kTfLiteFloat32:
      // 正常的全 float32 的运算
    case kTfLiteInt8:
      // wegiht 为 int8
      EvalQuantized<kernel_type>(context, node, params, data, input,
                                 filter, bias, output);

TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
                           TfLiteFullyConnectedParams* params, OpData* data,
                           const TfLiteTensor* input,
                           const TfLiteTensor* filter, const TfLiteTensor* bias,
                           TfLiteTensor* output)
  // 动态范围量化
  if (input->type == kTfLiteFloat32)
    return EvalHybrid(context, node, params, data, input, filter, bias,
                      input_quantized, scaling_factors, accum_scratch, row_sums,
                      input_offsets, output);      
  else
    // 全整数
    reference_ops::FullyConnected(
        op_params, GetTensorShape(input), GetTensorData<uint8_t>(input),
        GetTensorShape(filter), GetTensorData<uint8_t>(filter),
        GetTensorShape(bias), GetTensorData<int32_t>(bias),
        GetTensorShape(output), GetTensorData<uint8_t>(output));


TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node,
                        TfLiteFullyConnectedParams* params, OpData* data,
                        const TfLiteTensor* input, const TfLiteTensor* filter,
                        const TfLiteTensor* bias, TfLiteTensor* input_quantized,
                        TfLiteTensor* scaling_factors,
                        TfLiteTensor* accum_scratch, TfLiteTensor* row_sums,
                        TfLiteTensor* input_offsets, TfLiteTensor* output)
  // Quantize input from float to uint8 + quantization params (scaling factor).
  float* scaling_factors_ptr = GetTensorData<float>(scaling_factors);
  int8_t* quant_data = GetTensorData<int8_t>(input_quantized);
  const int8_t* filter_data = GetTensorData<int8_t>(filter);
  const float* input_ptr = GetTensorData<float>(input);
  tensor_utils::BatchQuantizeFloats(
    input_ptr, batch_size, input_size, quant_data, scaling_factors_ptr,
    input_offset_ptr, params->asymmetric_quantize_inputs);
  for (int b = 0; b < batch_size; ++b) {
    // Incorporate scaling of the filter.
    // !!! 最终使用的 scale 是 input_scale * filter_scale
    scaling_factors_ptr[b] *= filter->params.scale;
  }

  // Compute output += weight * quantized_input
  int32_t* scratch = GetTensorData<int32_t>(accum_scratch);
  tensor_utils::MatrixBatchVectorMultiplyAccumulate(
    filter_data, num_units, input_size, quant_data, scaling_factors_ptr,
    batch_size, GetTensorData<float>(output), /*per_channel_scale=*/nullptr,
    input_offset_ptr, scratch, row_sums_ptr, &data->compute_row_sums,
    CpuBackendContext::GetFromContext(context));

  // 最终的 output 是 float
  tensor_utils::ApplyActivationToVector(
    GetTensorData<float>(output), batch_size * num_units, params->activation,
    GetTensorData<float>(output));

// 矩阵运算时输入是两个 int8 的矩阵, 但输出的结果是 float
void PortableMatrixBatchVectorMultiplyAccumulate(
    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
    const int8_t* __restrict__ vectors, const float* scaling_factors,
    int n_batch, float* __restrict__ result) {
    for (int batch = 0; batch < n_batch; ++batch, vectors += m_cols) {
        const float batch_scaling_factor = scaling_factors[batch];
        // Get the address of the first row.
        const int8_t* row_ptr = matrix;
        for (int row = 0; row < m_rows; ++row) {
            // Initialize the dot product sum for the row to 0.
            int32_t dotprod = 0;
            for (int col = 0; col < m_cols; ++col, ++row_ptr) {
                dotprod += (*row_ptr) * (vectors[col]);
            }  // for col
            *result += dotprod * batch_scaling_factor;
            ++result;
        }  // for row
    }    // for batch
}

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model("/tmp/pow")

X = tf.random.uniform([100, 1], minval=1, maxval=10.0)


def representative_dataset_gen():
    for i in range(10):
        yield [[X[i]]]


converter.representative_dataset = representative_dataset_gen

converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

tflite = converter.convert()
with open("/tmp/pow-q.tflite", "wb") as f:
    f.write(tflite)
    print("size of tflite-q:", len(tflite))