mirror of
https://github.com/wassname/keras-js.git
synced 2026-08-21 11:16:47 +08:00
82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
import sys
|
|
import os
|
|
import h5py
|
|
import numpy as np
|
|
import json
|
|
|
|
|
|
class Encoder(object):
|
|
"""Encoder class.
|
|
Weights are serialized sequentially from the Keras flattened_layers representation
|
|
into:
|
|
- `weights`: a binary string representing the raw data bytes in float32
|
|
of all weights, sequentially concatenated.
|
|
- `metadata`: a list containing the byte length and tensor shape,
|
|
so that the original tensors can be reconstructed
|
|
"""
|
|
|
|
def __init__(self, weights_hdf5_filepath):
|
|
if not weights_hdf5_filepath:
|
|
raise Exception('weights_hdf5_filepath must be defined.')
|
|
self.weights_hdf5_filepath = weights_hdf5_filepath
|
|
self.weights = b''
|
|
self.metadata = []
|
|
|
|
def serialize(self):
|
|
"""serialize method.
|
|
Strategy for extracting the weights is adapted from the
|
|
load_weights_from_hdf5_group method of the Container class:
|
|
see https://github.com/fchollet/keras/blob/master/keras/engine/topology.py#L2505-L2585
|
|
"""
|
|
hdf5_file = h5py.File(self.weights_hdf5_filepath, mode='r')
|
|
if 'layer_names' not in hdf5_file.attrs and 'model_weights' in hdf5_file:
|
|
f = hdf5_file['model_weights']
|
|
else:
|
|
f = hdf5_file
|
|
|
|
layer_names = [n.decode('utf8') for n in f.attrs['layer_names']]
|
|
offset = 0
|
|
for layer_name in layer_names:
|
|
g = f[layer_name]
|
|
weight_names = [n.decode('utf8') for n in g.attrs['weight_names']]
|
|
if len(weight_names):
|
|
for weight_name in weight_names:
|
|
meta = {}
|
|
meta['layer_name'] = layer_name
|
|
meta['weight_name'] = weight_name
|
|
weight_value = g[weight_name].value
|
|
bytearr = weight_value.astype(np.float32).tobytes()
|
|
self.weights += bytearr
|
|
meta['offset'] = offset
|
|
meta['length'] = len(bytearr) // 4
|
|
meta['shape'] = list(weight_value.shape)
|
|
meta['type'] = 'float32'
|
|
self.metadata.append(meta)
|
|
offset += len(bytearr)
|
|
|
|
hdf5_file.close()
|
|
|
|
def save(self):
|
|
"""Saves weights data (binary) and weights metadata (json)
|
|
"""
|
|
weights_filepath = '{}_weights.buf'.format(os.path.splitext(self.weights_hdf5_filepath)[0])
|
|
with open(weights_filepath, mode='wb') as f:
|
|
f.write(self.weights)
|
|
metadata_filepath = '{}_metadata.json'.format(os.path.splitext(self.weights_hdf5_filepath)[0])
|
|
with open(metadata_filepath, mode='w') as f:
|
|
json.dump(self.metadata, f)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
"""
|
|
Usage:
|
|
python encoder.py example.hdf5
|
|
|
|
Output:
|
|
- example_weights.buf
|
|
- example_metadata.json
|
|
"""
|
|
encoder = Encoder(*sys.argv[1:])
|
|
encoder.serialize()
|
|
encoder.save()
|