from .backend import np, tf, keras
from ._main import BaseModel
from ai4water.tf_attributes import ACTIVATION_LAYERS, LAYERS, tcn, MULTI_INPUT_LAYERS
from .nn_tools import get_add_call_args, get_call_args
[docs]class Model(BaseModel):
"""
Model class with Functional API and inherits from `BaseModel`.
For ML/non-Neural Network based models, there is no difference in functional
or sub-clsasing api. For DL/NN-based models, this class implements functional
api and differs from subclassing api in internal implementation of NN. This
class is usefull, if you want to use the functional API of keras to build
your own NN structure. In such as case you can construct your NN structure
by overwriting `add_layers`. Another advantage of this class is that sometimes,
model_subclsasing is not possible for example due to some bugs in tensorflow.
In such a case this class can be used. Otherwise all the features of ai4water
are available in this class as well.
Example:
>>>from ai4water.functional import Model
"""
[docs] def __init__(self, *args, **kwargs):
"""
Initializes and builds the NN/ML model.
"""
self._go_up = True
super().__init__(*args, **kwargs)
def_KModel = None
if keras is not None:
def_KModel = keras.models.Model
self.KModel = kwargs.get('KModel', def_KModel)
self.build()
@property
def api(self):
return 'functional'
@property
def KModel(self):
"""sets k_model.
In case when we want to customize the model such as for implementing custom
`train_step`, we can provide the customized model as input the this Model
class
"""
return self._k_model
@KModel.setter
def KModel(self, x):
self._k_model = x
@property
def weights(self):
"""Returns names of weights in model."""
_ws = []
for w in self._model.weights:
_ws.append(w.name)
return _ws
@property
def layers(self):
if self.category == "ML":
raise NotImplementedError
return self._model.layers
@property
def inputs(self):
if self.category == "ML":
raise NotImplementedError
return self._model.inputs
@property
def outputs(self):
if self.category == "ML":
raise NotImplementedError
return self._model.outputs
@property
def output_shape(self)->tuple:
if self.category == "ML":
raise NotImplementedError
return self._model.output_shape
@property
def trainable_weights(self):
if self.category == "ML":
raise NotImplementedError
return self._model.trainable_weights
@property
def layer_names(self):
_all_layers = []
if self.category == "ML":
return None
for layer in self._model.layers:
_all_layers.append(layer.name)
return _all_layers
@property
def num_input_layers(self) -> int:
if self.category != "DL":
return np.inf
else:
return len(self._model.inputs)
@property
def input_layer_names(self) -> list:
return [lyr.name.split(':')[0] for lyr in self._model.inputs]
@property
def layers_out_shapes(self) -> dict:
""" returns shapes of outputs from all layers in model as dictionary"""
shapes = {}
for lyr in self._model.layers:
shapes[lyr.name] = lyr.output_shape
return shapes
@property
def layers_in_shapes(self) -> dict:
""" returns the shapes of inputs to all layers"""
shapes = {}
for lyr in self._model.layers:
shapes[lyr.name] = lyr.input_shape
return shapes
@property
def fit_fn(self):
return self._model.fit
@property
def evaluate_fn(self):
return self._model.evaluate
@property
def predict_fn(self):
return self._model.predict
def count_params(self):
if self.category == "ML":
raise NotImplementedError
return self._model.count_params()
def _get_dummy_input_shape(self):
shape = ()
if self.config['backend'] == 'tensorflow' and self.category == "DL":
if isinstance(self.model_.inputs, list):
if len(self.model_.inputs) == 1:
shape = self.model_.inputs[0].shape
else:
shape = [inp.shape for inp in self.model_.inputs]
return shape
def first_layer_shape(self):
""" instead of tuple, returning a list so that it can be moified if needed"""
if self.num_input_layers > 1:
shapes = {}
for lyr in self._model.inputs:
shapes[lyr.name] = lyr.shape
return shapes
shape = []
for idx, d in enumerate(self._model.layers[0].input.shape):
if int(tf.__version__[0]) == 1:
if isinstance(d, tf.Dimension): # for tf 1.x
d = d.value
if idx == 0: # the first dimension must remain undefined so that the user may define batch_size
d = -1
shape.append(d)
return shape
[docs] def add_layers(self, layers_config: dict, inputs=None):
"""
Builds the NN from dictionary.
Arguments:
layers_config : wholse keys can be one of the following:
`config`: `dict`/lambda, Every layer must contain initializing
arguments as `config` dictionary. The `config` dictionary
for every layer can contain `name` key and its value must be
`str` type. If `name` key is not provided in the config,
the provided layer name will be used as its name e.g in following case
layers = {'LSTM': {'config': {'units': 16}}}
the name of `LSTM` layer will be `LSTM` while in follwoing case
layers = {'LSTM': {'config': {'units': 16, 'name': 'MyLSTM'}}}
the name of the lstm will be `MyLSTM`.
`inputs`: str/list, The calling arguments for the list. If `inputs`
key is missing for a layer, it will be supposed that either
this is an Input layer or it uses previous outputs as inputs.
`outputs`: str/list We can specifity the outputs from a layer
by using the `outputs` key. The value to `outputs` must be a
string or list of strings specifying the name of outputs from
current layer which can be used later in the mdoel.
`call_args`: str/list We can also specify additional call arguments
by `call_args` key. The value to `call_args` must be a string
or a list of strings.
inputs : if None, it will be supposed the the `Input` layer either
exists in `layers_config` or an Input layer will be created
within this method before adding any other layer. If not None,
then it must be in `Input` layer and the remaining NN architecture
will be built as defined in `layers_config`. This can be handy
when we want to use this method several times to build a complex
or parallel NN structure. avoid `Input` in layer names.
Returns:
inputs :
outputs :
"""
lyr_cache = {}
wrp_layer = None # indicator for wrapper layers
first_layer = True
idx = 0
for lyr, lyr_args in layers_config.items():
idx += 0
if callable(lyr) and hasattr(lyr, '__call__'):
LAYERS[lyr.__name__] = lyr
self.config['model']['layers'] = update_layers_config(layers_config, lyr)
lyr = lyr.__name__
lyr_config, lyr_inputs, named_outs, call_args = self.deconstruct_lyr_args(lyr, lyr_args)
if callable(lyr) and not hasattr(lyr, '__call__'):
lyr = "lambda"
lyr_name, args, lyr_config, activation = self.check_lyr_config(lyr, lyr_config)
# may be user has defined layers without input layer, in this case add Input layer as first layer
if first_layer:
if inputs is not None: # This method was called by providing it inputs.
assert isinstance(inputs, tf.Tensor)
lyr_cache["Input"] = inputs
# since inputs have been defined, all the layers that will be added will be next to first layer
first_layer = False
layer_outputs = inputs
assign_dummy_name(layer_outputs, 'input')
elif lyr_name != "Input":
if 'input_shape' in lyr_config: # input_shape is given in the first layer so make input layer
layer_outputs = LAYERS["Input"](shape=lyr_config['input_shape'])
assign_dummy_name(layer_outputs, 'input')
else:
# for simple dense layer based models, lookback will not be used
def_shape = (self.num_ins,) if self.lookback == 1 else (self.lookback, self.num_ins)
layer_outputs = LAYERS["Input"](shape=def_shape)
# first layer is built so next iterations will not be for first layer
first_layer = False
# put the first layer in memory to be used for model compilation
lyr_cache["Input"] = layer_outputs
# add th layer which the user had specified as first layer
assign_dummy_name(layer_outputs, 'input')
if lyr_inputs is None: # The inputs to the layer have not been specified, so either it is an Input layer
# or it uses the previous outputs as inputs
if lyr_name == "Input":
# it is an Input layer, hence should not be called
layer_outputs = LAYERS[lyr_name](*args, **lyr_config)
assign_dummy_name(layer_outputs, 'input')
else:
# it is executable and uses previous outputs as inputs
if lyr_name in ACTIVATION_LAYERS:
layer_outputs = ACTIVATION_LAYERS[lyr_name](name=lyr_config['name'])(layer_outputs)
elif lyr_name in ['TimeDistributed', 'Bidirectional']:
wrp_layer = LAYERS[lyr_name]
lyr_cache[lyr_name] = wrp_layer
continue
elif "LAMBDA" in lyr_name.upper():
# lyr_config is serialized lambda layer, which needs to be deserialized
# by default the lambda layer takes the previous output as input
# however when `call_args` are provided, they overwrite the layer_outputs
if call_args is not None: # todo, add example in docs
layer_outputs = get_add_call_args(call_args, lyr_cache, lyr_config['name'])
layer_outputs = tf.keras.layers.deserialize(lyr_config)(layer_outputs)
# layers_config['lambda']['config'] still contails lambda, so we need to replace the python
# object (lambda) with the serialized version (lyr_config) so that it can be saved as json file.
layers_config[lyr]['config'] = lyr_config
else:
if wrp_layer is not None:
layer_outputs = wrp_layer(LAYERS[lyr_name](*args, **lyr_config))(layer_outputs)
wrp_layer = None
else:
add_args = get_add_call_args(call_args, lyr_cache, lyr_config['name'])
layer_initialized = LAYERS[lyr_name](*args, **lyr_config)
layer_outputs = layer_initialized(layer_outputs, **add_args)
self.get_and_set_attrs(layer_initialized)
else: # The inputs to this layer have been specified so they must exist in lyr_cache.
# it is an executable
if lyr_name in ACTIVATION_LAYERS:
call_args, add_args = get_call_args(lyr_inputs, lyr_cache, call_args, lyr_config['name'])
layer_outputs = ACTIVATION_LAYERS[lyr_name](name=lyr_config['name'])(call_args, **add_args)
elif lyr_name in ['TimeDistributed', 'Bidirectional']:
wrp_layer = LAYERS[lyr_name]
lyr_cache[lyr_name] = wrp_layer
continue
elif "LAMBDA" in lyr_name.upper():
call_args, add_args = get_call_args(lyr_inputs, lyr_cache, call_args, lyr_config['name'])
layer_outputs = tf.keras.layers.deserialize(lyr_config)(call_args)
layers_config[lyr]['config'] = lyr_config
else:
if wrp_layer is not None:
call_args, add_args = get_call_args(lyr_inputs, lyr_cache, call_args, lyr_config['name'])
layer_outputs = wrp_layer(LAYERS[lyr_name](*args, **lyr_config))(call_args, **add_args)
wrp_layer = None
else:
call_args, add_args = get_call_args(lyr_inputs, lyr_cache, call_args, lyr_config['name'])
layer_initialized = LAYERS[lyr_name](*args, **lyr_config)
# for multi-input layers inputs should be ([a,b,c]) instaed of (a,b,c)
if isinstance(lyr_inputs, list) and lyr_name in MULTI_INPUT_LAYERS:
layer_outputs = layer_initialized(*call_args, **add_args)
else:
layer_outputs = layer_initialized(call_args, **add_args)
self.get_and_set_attrs(layer_initialized)
if activation is not None: # put the string back to dictionary to be saved in config file
lyr_config['activation'] = activation
if named_outs is not None:
if isinstance(named_outs, (list, tuple)) or named_outs.__class__.__name__ in ["ListWrapper"]:
# this layer is returning more than one output
assert len(named_outs) == len(layer_outputs), "Layer {} is expected to return {} " \
"outputs but it actually returns " \
"{}".format(lyr_name, named_outs, layer_outputs)
for idx, out_name in enumerate(named_outs):
self.update_cache(lyr_cache, out_name, layer_outputs[idx])
else:
# this layer returns just one output, TODO, this might be re
self.update_cache(lyr_cache, named_outs, layer_outputs)
self.update_cache(lyr_cache, lyr_config['name'], layer_outputs)
first_layer = False
self.jsonize_lyr_config(lyr_config)
inputs = []
for k, v in lyr_cache.items():
# since the model is not build yet and we have access to only output tensors of each list, this is probably
# the only way to know that how many `Input` layers were encountered during the run of this method. Each
# tensor (except TimeDistributed) has .op.inputs attribute,
# which is empty if a tensor represents output of Input layer.
if int(''.join(tf.__version__.split('.')[0:2]).ljust(3, '0')) < 240:
if k != "TimeDistributed" and hasattr(v, 'op'):
if hasattr(v.op, 'inputs'):
_ins = v.op.inputs
if len(_ins) == 0:
inputs.append(v)
# not sure if this is the proper way of checking if a layer receives an input or not!
else:
if hasattr(v, '__dummy_name'):
inputs.append(v)
# for case when {Input -> Dense, Input_1}, this method wrongly makes Input_1 as output so in such case use
# {Input_1, Input -> Dense }, thus it makes Dense as output and first 2 as inputs, so throwing warning
if int(''.join(tf.__version__.split('.')[0:2]).ljust(3, '0')) < 240:
if len(layer_outputs.op.inputs) < 1:
print("Warning: the output is of Input tensor class type")
else:
if 'op' not in dir(layer_outputs): # layer_outputs does not have `op`, which means it has no incoming node
print("Warning: the output is of Input tensor class type")
return inputs, layer_outputs
def compile(self, model_inputs, outputs, **compile_args):
k_model = self.KModel(inputs=model_inputs, outputs=outputs)
k_model.compile(loss=self.loss(), optimizer=self.get_optimizer(), metrics=self.get_metrics(), **compile_args)
if self.verbosity > 0:
k_model.summary()
if self.verbosity >= 0:
self.plot_model(k_model)
return k_model
def build(self, input_shape=None):
self.print_info()
if self.category == "DL":
if self.config.get('model', None) is None:
lyrs = None
else:
lyrs = self.config['model']['layers']
inputs, predictions = self.add_layers(lyrs)
self._model = self.compile(inputs, predictions)
self.info['model_parameters'] = int(self._model.count_params()) if self._model is not None else None
if self.verbosity > 0 and self.config['model'] is not None:
if 'tcn' in self.config['model']['layers']:
if int(''.join(tf.__version__.split('.')[0:2]).ljust(3, '0')) >= 250:
# tf >= 2.5 does not have _layers and tcn uses _layers
setattr(self._model, '_layers', self._model.layers)
tcn.tcn_full_summary(self._model, expand_residual_blocks=True)
else:
self.build_ml_model()
if not getattr(self, 'from_check_point', False) and self.verbosity>=0:
# fit may fail so better to save config before as well. This will be overwritten once the fit is complete
self.save_config()
self.update_info()
return
def loss_name(self):
if isinstance(self._model.loss, str):
return self._model.loss
elif hasattr(self._model.loss, 'name'):
return self._model.loss.name
else:
return self._model.loss.__name__
def update_layers_config(layers_config, lyr):
new_config = {}
for k, v in layers_config.items():
if k == lyr:
new_config[lyr.__name__] = v
else:
new_config[k] = v
return new_config
def assign_dummy_name(tensor, dummy_name):
if isinstance(tensor, list):
for t in tensor:
setattr(t, '__dummy_name', dummy_name)
else:
setattr(tensor, '__dummy_name', dummy_name)