def get_config(self) -> Dict:
""" Get configuration of the Multi-Gate Mixture of Experts """
base_config = super().get_config()
return {
**base_config, "base_layer":
layers.serialize(self.base_layer) if self.base_layer else None,
"task_layers": [layers.serialize(l) for l in self.task_layers],
"moe_layers": [layers.serialize(l) for l in self.moe_layers],
"moe_dropout":
self.moe_dropout,
"moe_dropout_rate":
self.moe_dropout_rate
}
def get_layer_info(self):
"""Return Layer info"""
info = []
if self.type=="Input":
info.append({})
info[-1]["shape"] = list(self.tf_layer.get_shape())
info.append({})
else:
info.append(layers.serialize(self.tf_layer))
try:
# if layer is initialized and has weights
info[-1]["avg_weight"] = str(np.average(self.tf_layer.get_weights()[0]))
info[-1]["avg_bias"] = str(np.average(self.tf_layer.get_weights()[1]))
info[-1]["avg_abs_weight"] = str(np.average(np.abs(self.tf_layer.get_weights()[0])))
info[-1]["avg_abs_bias"] = str(np.average(np.abs(self.tf_layer.get_weights()[1])))
except:
# if layer is not initialized it won't have weights
info[-1]["avg_weight"] = 0
info[-1]["avg_bias"] = 0
info[-1]["avg_abs_weight"] = 0
info[-1]["avg_abs_bias"] = 0
info.append({})
info[-1]["id"] = self.name
info[-1]["type"] = self.type
info[-1]["Inputs"] = json.dumps(list(self.inp.keys()))
info[-1]["Outputs"] = json.dumps(list(self.out.keys()))
return info
def change_nodes(self, new_nodes):
# rebuild layer with 5 less nodes
# new_nodes = len(self.weights[1])-number_of_nodes
serialized = layers.serialize(self.tf_layer)
serialized["config"]["units"] = new_nodes
self.tf_layer = layers.deserialize(serialized)
# load back old weights
self.save_weights([self.weights[0][:,:new_nodes],
self.weights[1][:new_nodes]])
def parse_model(self, model):
"""Reads in a tensorflow model"""
# first clear all variables
self.layers = {}
self.connected_layers = {}
self.connected_inputs = []
self.connected_outputs = []
self.inputs = []
self.outputs = []
# iterate through all layers of the tensorflow model
for layer in model.layers:
# disconnected_layer = layers.deserialize({'class_name': layer.__class__.__name__, 'config': layer.get_config()})
disconnected_layer = layers.deserialize(layers.serialize(layer))
weights = layer.get_weights()
# add the layer to the self.layers array
self.layers[layer.name] = Layer(disconnected_layer, name = layer.name, type = "Layer")
# save the weights to the layer (NOTE this doesn't mean they're loaded to
# the tensorflow layer, they're only saved to the custom layer object)
self.layers[layer.name].weights = weights
# if there is only a single or no input make it into a list of length 1
# so it's iterable
if isinstance(layer._inbound_nodes[0].inbound_layers,list):
input_layer = layer._inbound_nodes[0].inbound_layers
else:
input_layer = [layer._inbound_nodes[0].inbound_layers]
# if the input_layer is empty we know that the layer has to be an
# input
if len(input_layer)<1:
self.layers[layer.name] = Layer(layer.input, name = layer.name, type = "Input")
self.inputs.append(layer.name)
# if layer.name not in self.layers:
# Otherwise we iterate though the inputs and set them accordingly
for input in input_layer:
self.set_input(input.name,layer.name)
# print(layer.name,"gets input from",input.name)
self.connect_layers()
for layer in self.layers:
# in_conv = False
if layer not in list(self.connections_out.keys()):
_, in_conv = self.find_value(layer,self.connections_in)
if not in_conv:
self.layers[layer].type = "Output"
self.outputs.append(self.layers[layer])
from tensorflow.keras import layers as kl
from tensorflow import keras
from tensorflow.keras.layers import Layer, Dense
import tensorflow as tf
from tensorflow.keras import regularizers
#%%
l = kl.Dense(10, kernel_regularizer="l2")
kl.serialize(l)
#%%
l1 = kl.deserialize({'class_name': 'Dense', 'config': {'units': 5}})
l1.get_config()
#%%
from tensorflow.python.keras.utils.generic_utils import deserialize_keras_object
deserialize_keras_object({
'class_name': 'Dense',
'config': {
'units': 5
}
},
module_objects=globals())
#%%
class ConstantMultiple(Layer):
def __init__(self, init_val: float = 1, regularizer=None, **kwargs):
super().__init__(**kwargs)
self.init_val = init_val
self.regularizer = regularizers.get(regularizer)
self.c = self.add_weight(name="c", shape=(), regularizer=regularizer)
def from_layer(cls,
layer,
name=None,
weights=None,
input_layers=None,
metadata=None):
"""Construct a LayerNode from a keras.layers.Layer object.
Args:
layer: keras.layers.Layer object.
name: string of layer name, layer.name will be used if it is not set.
weights: An OrderedDict of weight name => value for the layer, layer's weights
will be used if it is not set.
input_layers: List of `LayerNode`s that feed into this layer.
metadata: Dictionary of metadata for a given layer.
Return:
"""
layer_config = layers.serialize(layer)
if not name:
layer_config['name'] = layer.name
else:
layer_config['name'] = name
def _is_custom_layer(layer):
if layer is None:
return False
real_custom_objects = {"CustomOpWrapper", "Vitis>CustomOpWrapper"}
custom_objects = tf.keras.utils.get_custom_objects()
vitis_objects = set()
for cls in custom_objects:
if cls.startswith("Vitis>"):
vitis_objects.add(cls)
vitis_objects.add(cls.lstrip("Vitis>"))
# remove "CustomOpWrapper", "Vitis>CustomOpWrapper" from vitis_objects
vitis_objects = vitis_objects - real_custom_objects
for cls in custom_objects:
if cls not in vitis_objects:
real_custom_objects.add(cls)
if isinstance(layer, dict):
cls_name = layer['class_name']
else:
cls_name = layer.__class__.__name__
return (cls_name in real_custom_objects)
def _weight_name(name, keras_layer):
"""Extracts the weight name by removing layer from TF variable name.
For example, returns 'kernel:0' for 'dense_2/kernel:0'.
Args:
name: TensorFlow variable name.
layer: a keras layer or a layer config dict
Returns:
Extracted weight name.
"""
if _is_custom_layer(keras_layer):
return name
else:
return name.split('/')[-1]
def _get_keras_layer_weights(keras_layer):
"""Returns a map of weight name, weight matrix. Keeps keras ordering."""
weights_map = collections.OrderedDict()
for weight_tensor, weight_numpy in \
zip(keras_layer.weights, keras_layer.get_weights()):
weights_map[_weight_name(weight_tensor.name,
keras_layer)] = weight_numpy
return weights_map
if not weights:
weights = _get_keras_layer_weights(layer)
return cls(layer_config, weights, input_layers, metadata)
def __reduce__(self):
model_metadata = saving_utils.model_metadata(self)
training_config = model_metadata.get("training_config", None)
model = serialize(self)
weights = self.get_weights()
return (unpack, (model, training_config, weights))
def reset_layers(self):
for layer in self.layers:
if self.layers[layer].type != "Input":
self.layers[layer].tf_layer = layers.deserialize(layers.serialize(self.layers[layer].tf_layer))
else:
self.layers[layer].tf_layer._name = self.layers[layer].name