`quantize weights` C++ Examples

Status QuantizeWeights(const GraphDef& input_graph_def,
                       const TransformFuncContext& context,
                       GraphDef* output_graph_def) {
  int32 minimum_size;
  TF_RETURN_IF_ERROR(
      context.GetOneInt32Parameter("minimum_size", 1024, &minimum_size));
  TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
      input_graph_def, {"Const"},
      [minimum_size](const NodeMatch& match,
                     const std::set<string>& input_nodes,
                     const std::set<string>& output_nodes,
                     std::vector<NodeDef>* new_nodes) {
        const NodeDef& old_const_node = match.node;
        if (!old_const_node.attr().count("dtype")) {
          return errors::InvalidArgument("No 'dtype' attribute for Const node ",
                                         old_const_node.name());
        }
        if (!old_const_node.attr().count("value")) {
          return errors::InvalidArgument("No 'value' attribute for Const node ",
                                         old_const_node.name());
        }
        const DataType old_dtype = old_const_node.attr().at("dtype").type();
        Tensor old_tensor;
        if (!old_tensor.FromProto(old_const_node.attr().at("value").tensor())) {
          return errors::InvalidArgument("Decoding Tensor failed for node",
                                         old_const_node.name());
        }
        const size_t num_elements = old_tensor.NumElements();
        // If this isn't a float constant, or it's too small, then reuse the
        // same node with no changes.
        if ((old_dtype != DT_FLOAT) || (num_elements < minimum_size)) {
          new_nodes->push_back(old_const_node);
          return Status::OK();
        }
        const float* old_values = old_tensor.flat<float>().data();
        float min = std::numeric_limits<float>::max();
        float max = std::numeric_limits<float>::min();
        for (int i = 0; i < num_elements; ++i) {
          const float value = old_values[i];
          min = std::min(min, value);
          max = std::max(max, value);
        }
        // Make sure the quantization range includes 0.0f. Not all quantized
        // Ops behave properly if 0.0f is not in the range.
        min = std::min(min, 0.0f);
        max = std::max(0.0f, max);
        // min_value == max_value is a tricky case. It can occur for general
        // tensors, and of course for scalars. The quantized ops cannot deal
        // with this case, so we set max_value to something else.
        // It's a tricky question what is the numerically best solution to
        // deal with this degeneracy.
        // TODO(petewarden): Better use a tolerance than a hard comparison?
        if (min == max) {
          if (std::abs(min) < 0.000001f) {
            max = min + 1.0f;
          } else if (min > 0) {
            max = 2.0f * min;
          } else {
            max = min / 2.0f;
          }
        } 

auto quantize_weights(quantizer &quantizer, fake_kpu_conv2d &conv)
{
    auto weights = conv.weights();
    xt::xtensor<uint8_t, 4> q_weights(conv.weights().shape());
    std::vector<float> scales(conv.output_channels());
    auto total_range = quantizer.fixup_range(quantizer.get_range(weights.begin(), weights.end()));
    for (size_t oc = 0; oc < conv.output_channels(); oc++)
    {
        auto w_ch = xt::view(weights, oc, xt::all());
        auto range = quantizer.fixup_range(quantizer.get_range(w_ch.begin(), w_ch.end()));

        auto s1 = total_range.max / range.max;
        auto s2 = total_range.min / range.min;
        auto s = (s1 < 0 || s2 < 0) ? std::max(s1, s2) : std::min(s1, s2);

        assert(s > 0);
        for (auto &v : w_ch)
            v *= s;
        scales[oc] = s;
    }

    total_range = quantizer.get_range(weights.begin(), weights.end());
    auto q_p = quantizer.get_quant_param(total_range, 8);

    auto out_it = q_weights.begin();
    for (auto w : weights)
        *out_it++ = (uint8_t)std::clamp((int32_t)std::round(w * q_p.scale + q_p.zero_point), 0, 255);
    return std::make_tuple(q_p, std::move(scales), std::move(q_weights));
} 

auto quantize_weights(quantizer &quantizer, conv2d &conv)
{
    auto weights = conv.weights();
    xt::xtensor<uint8_t, 4> q_weights(conv.weights().shape());
    auto total_range = quantizer.fixup_range(quantizer.get_range(weights.begin(), weights.end()));
    auto q_p = quantizer.get_quant_param(total_range, 8);

    auto out_it = q_weights.begin();
    for (auto w : weights)
        *out_it++ = (uint8_t)std::clamp((int32_t)std::round(w * q_p.scale + q_p.zero_point), 0, 255);
    return std::make_tuple(q_p, std::move(q_weights));
}