Function decode_weights

tfjs-converter/python/tensorflowjs/read_weights.py:126–202

Load weight values from buffer(s) according to a weights manifest. Args: weights_manifest: A TensorFlow.js-format weights manifest (a JSON array). data_buffers: A buffer or a `list` of buffers containing the weights values in binary format, concatenated in the order specified in

(weights_manifest, data_buffers, flatten=False)

      offset=offset).reshape(shape)

def decode_weights(weights_manifest, data_buffers, flatten=False):
  """Load weight values from buffer(s) according to a weights manifest.

  Args:
    weights_manifest: A TensorFlow.js-format weights manifest (a JSON array).
    data_buffers: A buffer or a `list` of buffers containing the weights values
      in binary format, concatenated in the order specified in
      `weights_manifest`. If a `list` of buffers, the length of the `list`
      must match the length of `weights_manifest`. A single buffer is
      interpreted as a `list` of one buffer and is valid only if the length of
      `weights_manifest` is `1`.
    flatten: Whether all the weight groups in the return value are to be
      flattened as a single weight groups. Default: `False`.

  Returns:
    If `flatten` is `False`, a `list` of weight groups. Each group is an array
    of weight entries. Each entry is a dict that maps a unique name to a numpy
    array, for example:
        entry = {
          'name': 'weight1',
          'data': np.array([1, 2, 3], 'float32')
        }

        Weights groups would then look like:
        weight_groups = [
          [group_0_entry1, group_0_entry2],
          [group_1_entry1, group_1_entry2],
        ]
    If `flatten` is `True`, returns a single weight group.

  Raises:
    ValueError: if the lengths of `weights_manifest` and `data_buffers` do not
      match.
  """
  if not isinstance(data_buffers, list):
    data_buffers = [data_buffers]
  if len(weights_manifest) != len(data_buffers):
    raise ValueError(
        'Mismatch in the length of weights_manifest (%d) and the length of '
        'data buffers (%d)' % (len(weights_manifest), len(data_buffers)))

  out = []
  for group, data_buffer in zip(weights_manifest, data_buffers):
    offset = 0
    out_group = []

    for weight in group['weights']:
      quant_info = weight.get('quantization', None)
      name = weight['name']
      if weight['dtype'] == 'string':
        # String array.
        dtype = object
      elif quant_info:
        # Quantized array.
        dtype = np.dtype(quant_info['dtype'])
      else:
        # Regular numeric array.
        dtype = np.dtype(weight['dtype'])