Spaces:
Sleeping
Sleeping
File size: 4,202 Bytes
adcc0ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
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
|