Function ConvertToInt8Value

tensorflow/c/quantize_embedding_variable.cc:91–120

}

void ConvertToInt8Value(const Tensor& in_tensor, const string name,
                        const string scale_name, BundleWriter& writer) {
  auto in_data = in_tensor.flat<float>();
  TensorShape shape = in_tensor.shape();
  Tensor out_tensor(DataTypeToEnum<int8_t>::v(), shape);
  auto out_data = out_tensor.flat<int8_t>();
  int embed_dim = shape.dim_size(shape.dims() - 1);
  Tensor scale_tensor(DT_FLOAT, TensorShape({embed_dim}));
  auto scale_data = scale_tensor.flat<float>();
  std::vector<float> max_val(embed_dim, 0.0);
#if INTEL_MKL
#pragma omp parallel for num_threads(omp_get_num_procs())
#endif  // INTEL_MKL
  for (size_t i = 0; i < out_tensor.NumElements(); ++i) {
    int embed_i = i % embed_dim;
    max_val[embed_i] = std::max(max_val[embed_i], std::abs(in_data(i)));
  }
  for (size_t i = 0; i < embed_dim; ++i) {
    scale_data(i) = max_val[i] / 127.0;
  }
#if INTEL_MKL
#pragma omp parallel for num_threads(omp_get_num_procs())
#endif  // INTEL_MKL
  for (size_t i = 0; i < out_tensor.NumElements(); ++i) {
    int embed_i = i % embed_dim;
    out_data(i) = static_cast<int8_t>(round(in_data(i) / scale_data(embed_i)));
  }
  writer.Add(scale_name, scale_tensor);
  writer.Add(name, out_tensor);
}

Status QuantizeEmbeddingVariable(const string& input_prefix,
                                 const string& output_prefix,