Save keras model weights directly to bytes/memory?

0

Issue

Keras allows for saving entire models or just model weights (see thread). When saving the weights, they must be saved to a file, eg:

model = keras_model()
model.save_weights('/tmp/model.h5')

Instead of writing to file, I’d like to just save the bytes into memory. Something like

model.dump_weights()

Tensorflow doesn’t seem to have this, so as a workaround I’m writing to disk and then reading into memory:

temp = '/tmp/weights.h5'
model.save_weights(temp)
with open(temp, 'rb') as f:
    weightbytes = f.read()

Any way to avoid this roundabout?

Solution

Thanks @ddoGas for pointing out the model.get_weights() method, which returns a list of weights that can then be serialized. Just some context for why I am not saving the model in the conventional way: we are working with model wrapper classes that associate a model and custom behavior. For example, before prediction occurs special validation is needed:

class CNN:
   ...
   def predict():
       self.do_special_validation()
       self.model.predict()

Hence, we’re serializing the CNN class not just the underlying model. This is the solution to pickle the entire object. (pickle(CNN()) fails, otherwise we’d just use that)

import pickle

def serialize(cnn):
    return pickle.dumps({
        "weights": cnn.model.get_weights(),
        "cnnclass": cnn.__class__
    })

def deserialize(cnn_bytes):
    loaded = pickle.loads(cnn_bytes)
    weights, cnnclass = loaded['weights'], loaded['cnnclass']
    cnninstance = cnnclass()
    cnninstance.model.set_weights(weights)
    return cnninstance

Works well, thanks!

PS note using cnn.__class__ because don’t want to necessarily bind this to the CNN class directly but for it to work in general for any class that has a cnn.model attribute.

Answered By – Adam Hughes

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave A Reply

Your email address will not be published.

This website uses cookies to improve your experience. We'll assume you're ok with this, but you can opt-out if you wish. Accept Read More