Function read_weights

tfjs-converter/python/tensorflowjs/read_weights.py:35–75

Load weight values according to a TensorFlow.js weights manifest. Args: weights_manifest: A TensorFlow.js-format weights manifest (a JSON array). base_path: Base path prefix for the weights files. flatten: Whether all the weight groups in the return value are to be flattened as

(weights_manifest, base_path, flatten=False)

STRING_LENGTH_DTYPE = np.dtype('uint32').newbyteorder('<')

def read_weights(weights_manifest, base_path, flatten=False):
  """Load weight values according to a TensorFlow.js weights manifest.

  Args:
    weights_manifest: A TensorFlow.js-format weights manifest (a JSON array).
    base_path: Base path prefix for the weights files.
    flatten: Whether all the weight groups in the return value are to be
      flattened as a single weights group. 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.
  """
  if not isinstance(weights_manifest, list):
    raise ValueError(
        'weights_manifest should be a `list`, but received %s' %
        type(weights_manifest))

  data_buffers = []
  for group in weights_manifest:
    buff = io.BytesIO()
    buff_writer = io.BufferedWriter(buff)
    for path in group['paths']:
      with open(os.path.join(base_path, path), 'rb') as f:
        buff_writer.write(f.read())
    buff_writer.flush()
    buff_writer.seek(0)
    data_buffers.append(buff.read())
  return decode_weights(weights_manifest, data_buffers, flatten=flatten)


def _deserialize_string_array(data_buffer, offset, shape):