Add model argument to KerasModelWrapper
@@ -115,18 +115,16 @@ class KerasModelWrapper(Model):
in-place operations can incur an overhead.
"""
- def __init__(self):
+ def __init__(self, model=None):
"""
Create a wrapper for a Keras model
+ :param model: A Keras model
"""
super(KerasModelWrapper, self).__init__()
- self.model = None
- def set_model(self, model):
- """
- Store the Keras model being wrapped as an attribute.
- :param model: A Keras model
- """
+ if model is None:
+ raise ValueError('model argument must be suplied.')
+
self.model = model
def set_state(self, train):
@@ -145,10 +143,6 @@ class KerasModelWrapper(Model):
Looks for the name of the softmax layer.
:return: Softmax layer name
"""
- if self.model is None:
- raise ValueError('Call set_model() first with the Keras model '
- 'being wrapped.')
-
for i, layer in enumerate(self.model.layers):
cfg = layer.get_config()
if 'activation' in cfg and cfg['activation'] == 'softmax':
@@ -161,10 +155,6 @@ class KerasModelWrapper(Model):
Looks for the name of the layer producing the logits.
:return: name of layer producing the logits
"""
- if self.model is None:
- raise ValueError('Call set_model() first with the Keras model '
- 'being wrapped.')
-
softmax_name = self._get_softmax_name()
softmax_layer = self.model.get_layer(softmax_name)
node = softmax_layer.inbound_nodes[0]
@@ -177,10 +167,6 @@ class KerasModelWrapper(Model):
:param x: A symbolic representation of the network input.
:return: A symbolic representation of the logits
"""
- if self.model is None:
- raise ValueError('Call set_model() first with the Keras model '
- 'being wrapped.')
-
logits_name = self._get_logits_name()
return self.get_layer(x, logits_name)
@@ -190,10 +176,6 @@ class KerasModelWrapper(Model):
:param x: A symbolic representation of the network input.
:return: A symbolic representation of the probs
"""
- if self.model is None:
- raise ValueError('Call set_model() first with the Keras model '
- 'being wrapped.')
-
name = self._get_softmax_name()
return self.get_layer(x, name)
@@ -202,10 +184,6 @@ class KerasModelWrapper(Model):
"""
:return: Names of all the layers kept by Keras
"""
- if self.model is None:
- raise ValueError('Call set_model() first with the Keras model '
- 'being wrapped.')
-
layer_names = [x.name for x in self.model.layers]
return layer_names
@@ -216,10 +194,6 @@ class KerasModelWrapper(Model):
:return: A dictionary mapping layer names to the symbolic
representation of their output.
"""
- if self.model is None:
- raise ValueError('Call set_model() first with the Keras model '
- 'being wrapped.')
-
if (x, self.state) in self.fprop_cache.keys():
return self.fprop_cache[(x, self.state)]
else:
