| 174 | ) |
| 175 | |
| 176 | def get_opt_params(): |
| 177 | |
| 178 | trainable_modules = args.experiment.get('trainable_modules', None) |
| 179 | |
| 180 | # First, set all parameters to not require gradients |
| 181 | for param in model.parameters(): |
| 182 | param.requires_grad = False |
| 183 | |
| 184 | # If specific modules are provided, enable gradients for those |
| 185 | if trainable_modules: |
| 186 | # Enable gradients for specified modules |
| 187 | for name, module in model.named_parameters(): |
| 188 | if any(m in name for m in trainable_modules): |
| 189 | logger.info(f'Enabling gradients for {name}') |
| 190 | module.requires_grad = True |
| 191 | else: |
| 192 | # Enable gradients for all parameters |
| 193 | for param in model.parameters(): |
| 194 | param.requires_grad = True |
| 195 | |
| 196 | # Filter and return only trainable parameters |
| 197 | params_to_optimize = [p for p in model.parameters() if p.requires_grad] |
| 198 | logger.info(f'Trainable parameters: {sum([p.numel() for p in params_to_optimize])}') |
| 199 | return params_to_optimize |
| 200 | |
| 201 | if is_training: |
| 202 | opt = torch.optim.AdamW( |