diff --git a/torch_dreams/maco/features_visualizations/objectives.py b/torch_dreams/maco/features_visualizations/objectives.py index d942aa5..6330909 100644 --- a/torch_dreams/maco/features_visualizations/objectives.py +++ b/torch_dreams/maco/features_visualizations/objectives.py @@ -103,52 +103,7 @@ def __rmul__(self, factor: float): - def compile(self) -> Tuple[nn.Module, Callable, List[str], Tuple]: - """ - Compile all the sub-objectives into one and return the objects - for the optimization process. - - Returns - ------- - model_reconfigured - Model with the outputs needed for the optimization. - objective_function - Function to call that compute the loss for the objectives. - names - Names of each objectives. - input_shape - Shape of the input, one sample for each optimization. - """ - # the number of inputs will be the number of combinations possible - # of the objectives, the mask are used to take into account - # these combinations - nb_sub_objectives = len(self.multipliers) - - # re-arrange to match the different objectives with the model outputs - masks = np.array([np.array(m, dtype=object) for m in itertools.product(*self.masks)]) - masks = [torch.tensor(masks[:, i], dtype=torch.float32) for i in range(nb_sub_objectives)] - - # the name of each combination is the concatenation of each objectives - names = np.array([' & '.join(names) for names in itertools.product(*self.names)]) - # one multiplier by sub-objective - multipliers = torch.tensor(self.multipliers) - - def objective_function(model_outputs): - loss = 0.0 - for output_index in range(0, nb_sub_objectives): - outputs = model_outputs[output_index] - loss += self.funcs[output_index]( - outputs, masks[output_index].to(outputs.device)) - loss *= multipliers[output_index] - return loss - - # the model outputs will be composed of the layers needed - model_reconfigured = nn.Sequential(*self.layers) - - nb_combinations = masks[0].shape[0] - input_shape = (nb_combinations, *model_reconfigured.input_shape[1:]) - - return model_reconfigured, objective_function, names, input_shape + @@ -195,9 +150,11 @@ def layer(model: nn.Module, power = 2.0 if reducer == "magnitude" else 1.0 - def optim_func(model_output, mask): - return torch.mean((model_output * mask) ** power) + def optim_func(model_output, mask, power): + result = (model_output * mask) ** power + return torch.mean(result) + return Objective(model, [layer_output], [mask], [optim_func], [multiplier], [name]) @@ -274,25 +231,26 @@ def optim_func(model_output, mask): @staticmethod def channel(model: nn.Module, layer: str, - vectors:Union[torch.tensor,List[torch.tensor]], + channel_ids: Union[int, List[int]], multiplier: float = 1.0, - cossim_pow: float = 2.0, names: Optional[Union[str, List[str]]] = None): """ - Util to build an objective to maximise a channel in a layer. + util to build an objective to maximise a channel in a layer. Parameters ---------- model Model used for optimization. + layer Index or name of the targeted layer. - vectors - List of vectors to optimize. + + channel_ids + List of channels to optimize. + multiplier Multiplication factor of the objective. - cos_sim_pow - Power of the cosine similarity. + names A name for the objective. @@ -302,47 +260,56 @@ def channel(model: nn.Module, An objective ready to be compiled """ - layer_output = extract_features(model, layer) layer_shape = get_layer_output_shape(model, layer) - + layer_output = extract_features(model, layer) channel_ids = channel_ids if isinstance(channel_ids, list) else [channel_ids] masks = np.zeros((len(channel_ids), *layer_shape[1:])) - for i, channel_id in enumerate(channel_ids): - masks[i, ..., channel_id] = 1.0 + for i,c_id in enumerate(channel_ids): + masks[i, ..., c_id] = 1.0 - if names is None: - names = [f"Channel#{layer}_{i}" for i in channel_ids] + layer_name = get_layer_name(layer) + if names is None: + name = [f"Channel#{layer_name}_{c_id}" for c_id in channel_ids] axis_to_reduce = list(range(1, len(layer_shape))) - def optim_func(output, target): - return torch.mean(output * target, axis=axis_to_reduce) - + def optim_func(output,target,axis_to_reduce = None): + product = output * target + if axis_to_reduce is not None: + return torch.mean(product, dim=axis_to_reduce) + else: + return torch.mean(product) + return Objective(model, [layer_output], [masks], [optim_func], [multiplier], [names]) - - + + @staticmethod - def neuron(model:nn.Module, - layer:str, - neuron_ids:Union[int,List[int]], - multiplier:float=1.0, - names:Optional[Union[str,List[str]]]=None): + def neuron(model: nn.Module, + layer: str, + neuron_ids: Union[int, List[int]], + multiplier: float = 1.0, + names: Optional[Union[str, List[str]]] = None): """ Util to build an objective to maximise a neuron in a layer. Parameters ---------- + model Model used for optimization. + layer Index or name of the targeted layer. + neuron_ids - List of neuron ids to optimize. + List of neurons to optimize. + multiplier - Multiplication factor of the objective. + Multiplication factor of the objectives. + names A name for the objective. @@ -353,65 +320,40 @@ def neuron(model:nn.Module, """ layer_output = extract_features(model, layer) + neurons_ids = neurons_ids if isinstance(neurons_ids, list) else [neurons_ids] + nb_objectives = len(neurons_ids) layer_shape = get_layer_output_shape(model, layer) - neuron_ids = neuron_ids if isinstance(neuron_ids, list) else [neuron_ids] - - nb_objectives = len(neuron_ids) - - - layer_shape = get_layer_output_shape(model, layer) layer_shape = layer_shape[1:] - masks = np.zeros((nb_objectives, *layer_shape)) - masks = masks.reshape((nb_objectives, -1)) - for i, neuron_id in enumerate(neuron_ids): + + for i, neuron_id in enumerate(neurons_ids): masks[i, neuron_id] = 1.0 masks = masks.reshape((nb_objectives, *layer_shape)) + layer_name = get_layer_name(layer) if names is None: - names = [f"Neuron#{layer}_{i}" for i in neuron_ids] + names = [f"Neuron#{layer_name}_{neuron_id}" for neuron_id in neurons_ids] axis_to_reduce = list(range(1, len(layer_shape)+1)) - def optim_func(output, target): - return torch.mean(output * target, axis=axis_to_reduce) - - return Objective(model, [layer_output], [masks], [optim_func], [multiplier], [names]) - - + def optim_func(output,target): + product = output * target + return torch.mean(product, dim=axis_to_reduce) + + + return Objective(model, [layer_output], [masks], [optim_func], [multiplier], [names]) - - - - - - - - - - - - - - - - - - - - - + - - + \ No newline at end of file