Spaces:
Sleeping
Sleeping
| import torch | |
| from typing import Dict | |
| from utils.regularizations import * | |
| from optimization.ActivationExtractor import * | |
| def optimize_activation_batch(model, | |
| target_neurons:list, | |
| input_data:torch.Tensor, | |
| target_layer:torch.nn.Module, | |
| lr:float=0.01, | |
| steps:int=50, | |
| regularizations_dict:Dict=None, # this dict is being declared in main.py | |
| log_freq: int=10): | |
| """ | |
| Optimize input to maximize activation of a specific neuron with dynamic regularizations. | |
| Args: | |
| model: Pre-trained model set to evaluation mode. | |
| target_neuron (list): list of the neurons to maximize. | |
| input_data (torch.Tensor): Input tensor to optimize. Must have requires_grad=True. | |
| lr (float): Learning rate for the optimizer. | |
| steps (int): Number of optimization steps. | |
| regularizations_dict (dict): Dictionary of regularization functions and their parameters. | |
| Supported keys: | |
| - 'l2': L2 regularization weight (float). | |
| - 'tv': Total Variation loss weight (float). | |
| - 'sparsity': Sparsity constraint weight (float). | |
| - 'clip': Tuple of (min_value, max_value) for clipping input values. | |
| log_freq (int): Frequency of logging optimization progress. | |
| reduction_method (str): Method to aggregate neuron activations ('mean' or 'sum'). | |
| Returns: | |
| dict: A dictionary containing: | |
| - 'optimized_input': The input tensor optimized to maximize the target neuron's activation. | |
| - 'loss_history': A list of loss values recorded during the optimization process. | |
| """ | |
| optimizer = torch.optim.Adam(params=[input_data], lr=lr) | |
| regularizations_dict = regularizations_dict or {} | |
| reg_functions = { | |
| 'l2': l2_regularization, | |
| 'tv': total_variance_loss, | |
| 'sparsity': sparsity_constraint, | |
| 'entropy':entropy_loss, | |
| 'l1': l1_regularization, | |
| 'linf':linf_regularization, | |
| 'feature_map_sparsity':feature_map_sparsity | |
| } | |
| results = {"optimized_input": input_data, "loss_history": []} | |
| with ActivationExtractor(model=model, target_layer=target_layer) as extractor: | |
| for step in range(steps): | |
| _ = model(input_data) | |
| activation = extractor.activation | |
| if not target_neurons or max(target_neurons) >= activation.size(1): | |
| raise ValueError(f"Invalid target neurons:{target_neurons}. Layer has {activation.size(1)} neurons.") | |
| if len(target_neurons)==1: | |
| loss = - activation[0, target_neurons[0]].sum() | |
| else: | |
| reduction_method = regularizations_dict.get("reduction", "mean") # is used to dynamically retrieve a key ("reduction") from the regularizations_dict dictionary, with a default value of "mean" if the key is not provided. | |
| if reduction_method == "mean": | |
| loss = - activation[0,target_neurons].mean() #making it scalar | |
| elif reduction_method == "sum": | |
| loss = - activation[0,target_neurons].sum() | |
| else: | |
| raise ValueError(f"Unspported Reduction Method: {reduction_method}. Supported methods are 'mean' and 'sum'.") | |
| #print(f"reduction:{reduction_method}") | |
| for reg_name, reg_weight in regularizations_dict.items(): | |
| if reg_name=='clip': | |
| continue | |
| elif reg_name in reg_functions: | |
| if reg_name == "feature_map_sparsity": | |
| loss += reg_functions[reg_name](activation = activation, weight = reg_weight) | |
| else: | |
| loss += reg_functions[reg_name](input_data = input_data, weight = reg_weight) | |
| optimizer.zero_grad() | |
| loss.backward() | |
| optimizer.step() | |
| if 'clip' in regularizations_dict: | |
| with torch.no_grad(): | |
| input_data.clamp(*regularizations_dict['clip']) | |
| results["loss_history"].append(loss.item()) | |
| if log_freq and step % log_freq == 0: | |
| print(f"Step {step}/{steps} | Loss: {loss.item():.4f} | Neurons: {target_neurons}") | |
| results['optimized_input'] = input_data | |
| return results | |