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", )