def test_delete_layer_same_layer_outputs():
# Create all model layers
input_1 = Input(shape=(10, ))
dense_1 = Dense(3)
dense_2 = Dense(3)
dense_3 = Dense(3)
dense_4 = Dense(1)
# Create the base model
x = dense_1(input_1)
y = dense_2(x)
x = dense_3(x)
output_1 = dense_4(x)
output_2 = dense_4(y)
model_1 = utils.clean_copy(Model(input_1, [output_1, output_2]))
# Create the expected modified model
x = dense_1(input_1)
y = dense_2(x)
output_1 = dense_4(x)
output_2 = dense_4(y)
model_2_exp = utils.clean_copy(Model(input_1, [output_1, output_2]))
# Delete layer dense_3
model_2 = operations.delete_layer(model_1,
model_1.get_layer(dense_3.name),
copy=False)
# Compare the modified model with the expected modified model
assert compare_models(model_2, model_2_exp)
def test_delete_all_channels_in_long_branch():
input_1 = Input(shape=(20, 20, 3))
conv_1 = Conv2D(2, [3, 3], name="conv_1")
conv_2 = Conv2D(3, [3, 3], name="conv_2")
conv_3 = Conv2D(4, [1, 1], name="conv_3")
cat_1 = Concatenate(name="cat_1")
x = conv_1(input_1)
x = conv_3(x)
y = conv_2(input_1)
output_1 = cat_1([x, y])
model_1 = utils.clean_copy(Model(input_1, output_1))
surgeon = Surgeon(model_1, copy=True)
surgeon.add_job("delete_channels",
model_1.get_layer("conv_1"),
channels=[0, 1])
model_2 = surgeon.operate()
output_1 = conv_2(input_1)
model_2_exp = utils.clean_copy(Model(input_1, output_1))
config_1 = model_2.get_config()
config_2 = model_2_exp.get_config()
config_2["name"] = config_1["name"] # make the config names identical
assert config_1 == config_2
def test_delete_layer():
# Create all model layers
input_1 = Input(shape=[7, 7, 1])
conv2d_1 = Conv2D(3, [3, 3], data_format="channels_last")
conv2d_2 = Conv2D(3, [3, 3], data_format="channels_last")
flatten_1 = Flatten()
dense_1 = Dense(3)
dense_2 = Dense(3)
dense_3 = Dense(3)
dense_4 = Dense(1)
# Create the base model
x = conv2d_1(input_1)
x = conv2d_2(x)
x = flatten_1(x)
x = dense_1(x)
x = dense_2(x)
x = dense_3(x)
output_1 = dense_4(x)
model_1 = utils.clean_copy(Model(input_1, output_1))
# Create the expected modified model
x = conv2d_1(input_1)
x = conv2d_2(x)
x = flatten_1(x)
x = dense_1(x)
x = dense_3(x)
output_2 = dense_4(x)
model_2_exp = utils.clean_copy(Model(input_1, output_2))
# Delete layer dense_2
model_2 = operations.delete_layer(model_1, model_1.get_layer(dense_2.name))
# Compare the modified model with the expected modified model
assert compare_models(model_2, model_2_exp)
def test_delete_channels_downstream_sharing():
# Create all model layers
input_1 = Input(shape=(5,))
dense_1 = Dense(4, name='dense_1')
dense_2 = Dense(4, name='dense_2')
dense_3 = Dense(3, name='dense_3')
# Create the base model
x = dense_1(input_1)
y = dense_2(input_1)
output_1 = dense_3(x)
output_2 = dense_3(y)
model_1 = utils.clean_copy(Model(input_1, [output_1, output_2]))
# Delete channels from dense_1 and dense_2
surgeon = Surgeon(model_1)
surgeon.add_job('delete_channels', model_1.get_layer(dense_1.name), channels=[0])
surgeon.add_job('delete_channels', model_1.get_layer(dense_2.name), channels=[1])
model_2 = surgeon.operate()
# Create the expected model
# input_1 = Input(shape=(5,))
dense_1_exp = Dense(3, name='dense_1')
dense_2_exp = Dense(3, name='dense_2')
dense_3_exp = Dense(3, name='dense_3')
# Create the base model
x = dense_1_exp(input_1)
y = dense_2_exp(input_1)
output_1 = dense_3_exp(x)
output_2 = dense_3_exp(y)
model_2_exp = utils.clean_copy(Model(input_1, [output_1, output_2]))
config_1 = model_2.get_config()
config_2 = model_2_exp.get_config()
config_2['name'] = config_1['name'] # make the config names identical
assert json.dumps(config_1) == json.dumps(config_2)
def operate(self):
"""Perform all jobs assigned to the surgeon.
"""
# Operate on each node in self.nodes by order of decreasing depth.
sorted_nodes = sorted(
self.nodes,
reverse=True,
key=lambda x: utils.get_node_depth(self.model, x))
for node in sorted_nodes:
# Rebuild submodel up to this node
sub_output_nodes = utils.get_node_inbound_nodes(node)
outputs, output_masks = self._rebuild_graph(
self.model.inputs, sub_output_nodes)
# Perform surgery at this node
kwargs = self._kwargs_map[node]
self._mod_func_map[node](node, outputs, output_masks, **kwargs)
# Finish rebuilding model
output_nodes = []
for output in self.model.outputs:
layer, node_index, tensor_index = output._keras_history
output_nodes.append(get_inbound_nodes(layer)[node_index])
new_outputs, _ = self._rebuild_graph(self.model.inputs, output_nodes)
new_model = self._model_cls(self.model.inputs, new_outputs)
if self._copy:
return utils.clean_copy(new_model, self._custom_objects)
else:
return new_model
def operate(self):
"""Perform all jobs assigned to the surgeon.
"""
# Operate on each node in self.nodes by order of decreasing depth.
sorted_nodes = sorted(
self.nodes,
reverse=True,
key=lambda x: utils.get_node_depth(self.model, x))
for node in sorted_nodes:
# Rebuild submodel up to this node
sub_output_nodes = utils.get__inbound_nodes(node)
outputs, output_masks = self._rebuild_graph(
self.model.inputs, sub_output_nodes)
# Perform surgery at this node
kwargs = self._kwargs_map[node]
self._mod_func_map[node](node, outputs, output_masks, **kwargs)
# Finish rebuilding model
output_nodes = [
self.model.output_layers[i]._inbound_nodes[node_index] for i,
node_index in enumerate(self.model.output_layers_node_indices)
]
new_outputs, _ = self._rebuild_graph(self.model.inputs, output_nodes)
new_model = Model(self.model.inputs, new_outputs)
if self._copy:
return utils.clean_copy(new_model)
else:
return new_model
def test_delete_all_channels_in_branch():
input_1 = Input(shape=(20, 20, 3))
conv_1 = Conv2D(2, [3, 3], name='conv_1')
conv_2 = Conv2D(3, [3, 3], name='conv_2')
cat_1 = Concatenate(name='cat_1')
x = conv_1(input_1)
y = conv_2(input_1)
output_1 = cat_1([x, y])
model_1 = utils.clean_copy(Model(input_1, output_1))
surgeon = Surgeon(model_1, copy=True)
surgeon.add_job('delete_channels', model_1.get_layer('conv_1'), channels=[0, 1])
model_2 = surgeon.operate()
output_1 = conv_2(input_1)
model_2_exp = utils.clean_copy(Model(input_1, output_1))
config_1 = model_2.get_config()
config_2 = model_2_exp.get_config()
config_2['name'] = config_1['name'] # make the config names identical
assert json.dumps(config_1) == json.dumps(config_2)
def test_insert_layer():
# Create all model layers
input_1 = Input(shape=[7, 7, 1])
dense_1 = Dense(3)
dense_2 = Dense(3)
dense_3 = Dense(3)
dense_4 = Dense(1)
# Create the model
x = dense_1(input_1)
x = dense_2(x)
output_1 = dense_4(x)
model_1 = utils.clean_copy(Model(input_1, output_1))
# Create the expected modified model
x = dense_1(input_1)
x = dense_2(x)
x = dense_3(x)
output_2 = dense_4(x)
model_2_exp = utils.clean_copy(Model(input_1, output_2))
# Insert dense_3 before dense_4 in model_1
model_2 = operations.insert_layer(model_1, model_1.get_layer(dense_4.name),
dense_3)
# Compare the modified model with the expected modified model
assert compare_models(model_2, model_2_exp)
def __init__(self, model, copy=None):
if copy:
self.model = utils.clean_copy(model)
else:
self.model = model
self.nodes = []
self._copy = copy
self._finished_nodes = {}
self._replace_tensors = {}
self._channels_map = {}
self._new_layers_map = {}
self._insert_layers_map = {}
self._replace_layers_map = {}
self._mod_func_map = {}
self._kwargs_map = {}
self.valid_jobs = ('delete_layer', 'insert_layer', 'replace_layer',
'delete_channels')
def __init__(self, model, copy=None, custom_objects=None):
self._custom_objects = custom_objects
if copy:
self.model = utils.clean_copy(model, self._custom_objects)
else:
self.model = model
self._model_cls = self.model.__class__
check_tf_model(self.model)
self.nodes = []
self._copy = copy
self._finished_nodes = {}
self._replace_tensors = {}
self._channels_map = {}
self._new_layers_map = {}
self._insert_layers_map = {}
self._replace_layers_map = {}
self._mod_func_map = {}
self._kwargs_map = {}
self.valid_jobs = (
"delete_layer",
"insert_layer",
"replace_layer",
"delete_channels",
)
