diff --git a/src/detoxai/methods/clarcs/rrclarc.py b/src/detoxai/methods/clarcs/rrclarc.py index 272bbea..0a92d22 100644 --- a/src/detoxai/methods/clarcs/rrclarc.py +++ b/src/detoxai/methods/clarcs/rrclarc.py @@ -75,14 +75,55 @@ def apply_model_correction( assert len(cav_layers) == 1, "RR-CLARC only supports one CAV layer" self.cav_layer = cav_layers[0] - # Register rr_clarc_hook + # Register rr_clarc_hook and build list of module names in order + cav_layer_found = False + module_names_in_order = [] for name, module in self.model.named_modules(): + if name: # Skip empty string (root module) + module_names_in_order.append(name) if name == self.cav_layer: hook_fn = self.rr_clarc_hook() - handle = module.register_forward_hook(hook_fn) self.hooks.append(handle) _logger.debug(f"Added RR-CLARC hook to layer: {name}") + cav_layer_found = True + # Continue to collect all module names for proper layer ordering + + if not cav_layer_found: + raise ValueError(f"CAV layer '{self.cav_layer}' not found in model") + + # Get the index of cav_layer and the set of modules to freeze (all before cav_layer) + cav_layer_idx = module_names_in_order.index(self.cav_layer) + modules_to_freeze = set(module_names_in_order[:cav_layer_idx]) + + # Store original requires_grad state and freeze layers before cav_layer + self.original_requires_grad = {} + for name, param in self.model.named_parameters(): + self.original_requires_grad[name] = param.requires_grad + + # Check if this parameter belongs to a module that should be frozen + # For parameter 'layer1.0.conv.weight', we check if 'layer1', 'layer1.0', + # or 'layer1.0.conv' is in modules_to_freeze + if "." not in name: + # Top-level parameters (unlikely in practice) - skip freezing + continue + + param_module = name.rsplit(".", 1)[ + 0 + ] # Extract module path from parameter name + should_freeze = False + + # Check all possible parent module paths (break early on first match) + parts = param_module.split(".") + for i in range(1, len(parts) + 1): + module_path = ".".join(parts[:i]) + if module_path in modules_to_freeze: + should_freeze = True + break # Found a match, no need to check further + + if should_freeze: + param.requires_grad = False + _logger.debug(f"Frozen parameter: {name}") # Override training_step in lightning model by modified_training_step clone_original_training_step = deepcopy(self.lightning_model.training_step) @@ -92,7 +133,9 @@ def apply_model_correction( def configure_optimizers(self): """ """ - optimizer = torch.optim.Adam(self.parameters(), lr=ft_lr) + # Only optimize parameters that require gradients (from cav_layer onwards) + params_to_optimize = [p for p in self.parameters() if p.requires_grad] + optimizer = torch.optim.Adam(params_to_optimize, lr=ft_lr) return optimizer self.lightning_model.configure_optimizers = types.MethodType( @@ -117,6 +160,11 @@ def configure_optimizers(self): # Go back to eval mode self.lightning_model.eval() + # Restore original requires_grad state + for name, param in self.model.named_parameters(): + if name in self.original_requires_grad: + param.requires_grad = self.original_requires_grad[name] + # Remove hooks self.remove_hooks()