Function quantize_weights

tfjs-converter/python/tensorflowjs/quantization.py:93–147

Quantizes the weights by linearly re-scaling across available bits. The weights are quantized by linearly re-scaling the values between the minimum and maximum value, and representing them with the number of bits provided by the `quantization_dtype`. In order to guarantee that 0 is perfect

(data, quantization_dtype)

  return quantization_dtype

def quantize_weights(data, quantization_dtype):
  """Quantizes the weights by linearly re-scaling across available bits.

  The weights are quantized by linearly re-scaling the values between the
  minimum and maximum value, and representing them with the number of bits
  provided by the `quantization_dtype`.

  In order to guarantee that 0 is perfectly represented by one of the quantized
  values, the range is "nudged" in the same manner as in TF-Lite.

  Weights can be de-quantized by multiplying by the returned `scale` and adding
  `min`.

  Args:
    data: A numpy array of dtype 'float32' or 'int32'.
    quantization_dtype: A numpy dtype to quantize weights to. Only np.float16,
      np.uint8, and np.uint16 are supported.

  Returns:
    quantized_data: The quantized weights as a numpy array with dtype
      `quantization_dtype`.
    metadata: A dictionary with the corresponding metadata for the quantization
      type. There is no metadata associated with float16.
      For affine quantization there are two associated metadata values:
        scale: The linearly scaling constant used for quantization.
        min_val: The minimum value of the linear range.
  Raises:
    ValueError: if `quantization_dtype` is not a valid type.
  """
  if quantization_dtype in [np.uint8, np.uint16]:
    # Compute the min and max for the group.
    min_val = data.min().astype(np.float64)
    max_val = data.max().astype(np.float64)
    if min_val == max_val:
      # If there is only a single value, we can represent everything as zeros.
      quantized_data = np.zeros_like(data, dtype=quantization_dtype)
      scale = 1.0
    else:
      # Quantize data.
      scale, min_val, max_val = _get_affine_quantization_range(
          min_val, max_val, quantization_dtype)
      quantized_data = np.round(
          (data.clip(min_val, max_val) - min_val) / scale).astype(
              quantization_dtype)

    return quantized_data, {'min': min_val, 'scale': scale}
  elif quantization_dtype == np.float16:
    if data.dtype != np.float32:
      raise ValueError(
          'Invalid data dtype %r\n'
          'float16 quantization only supports float32 dtype' % data.dtype)
    quantized_data = data.astype(np.float16)
    return quantized_data, {}
  else:
    raise ValueError('Invalid `quantization_dtype`: %r' % quantization_dtype)