data aware
(
args,
model,
model_base=None,
prune_n=0,
prune_m=0,
prune_mode="activation",
name_filter_fn=None,
)
| 319 | |
| 320 | |
| 321 | def _prune_core( |
| 322 | args, |
| 323 | model, |
| 324 | model_base=None, |
| 325 | prune_n=0, |
| 326 | prune_m=0, |
| 327 | prune_mode="activation", |
| 328 | name_filter_fn=None, |
| 329 | ): |
| 330 | """ |
| 331 | data aware |
| 332 | """ |
| 333 | assert not args.prune_part, "Warning: prune_part is not supported" |
| 334 | # assert not args.neg_prune, "Warning: neg_prune is not supported" |
| 335 | prune_data = args.prune_data |
| 336 | for name, module in model.named_modules(): |
| 337 | if name_filter_fn is not None and not name_filter_fn(name): |
| 338 | continue |
| 339 | |
| 340 | if isinstance(module, ActLinear): |
| 341 | print("pruning:", name) |
| 342 | |
| 343 | i = re.search(r"\d+", name) |
| 344 | if i: |
| 345 | i = int(i.group()) |
| 346 | else: |
| 347 | i = 0 |
| 348 | |
| 349 | print("layer", i) |
| 350 | |
| 351 | if model_base is not None: |
| 352 | module_base = model_base.get_submodule(name) |
| 353 | |
| 354 | if args.use_diff: |
| 355 | magnitude = torch.abs(module.base.weight.data - module_base.weight.data) |
| 356 | else: |
| 357 | magnitude = torch.abs(module.base.weight.data) |
| 358 | |
| 359 | if prune_mode == "activation": |
| 360 | act = (module.activation_norms**0.5).unsqueeze(0) |
| 361 | elif prune_mode == "gradient": |
| 362 | act = module.base.weight.grad.abs() |
| 363 | else: |
| 364 | raise NotImplemented |
| 365 | |
| 366 | W_metric = magnitude * act |
| 367 | if args.neg_prune: |
| 368 | W_metric = -W_metric |
| 369 | |
| 370 | # copied from lib/prune.py prune_wanda: |
| 371 | |
| 372 | if args.dump_wanda_score: |
| 373 | # Only save the score, no pruning |
| 374 | save_folder = os.path.join( |
| 375 | args.save, f"wanda_score/" |
| 376 | ) # We assume that args.save has contained the information of pruned data. |
| 377 | if not os.path.exists(save_folder): |
| 378 | os.makedirs(save_folder) |
no test coverage detected