perforatedai.utils_perforatedai
1# Copyright (c) 2025 Perforated AI 2 3import torch 4import torch.nn as nn 5import torch.nn.init as init 6import torch.nn.functional as F 7import math 8import sys 9import numpy as np 10import pdb 11import os 12import time 13import warnings 14from collections import defaultdict 15 16from perforatedai import globals_perforatedai as GPA 17from perforatedai import modules_perforatedai as PA 18from perforatedai import tracker_perforatedai as TPA 19from perforatedai import clean_perforatedai as CL 20from perforatedai import blockwise_perforatedai as BPA 21from perforatedai import network_perforatedai as NPA 22from perforatedai import configure_perforatedai as CPA 23 24try: 25 from dashboard_utils.event_emitter import emitter as _dashboard_emitter 26except ImportError: 27 _dashboard_emitter = None 28 29try: 30 from perforatedbp import utils_pbp as UPB 31 from perforatedbp import modules_pbp as MPB 32except ModuleNotFoundError as e: 33 # Only pass if perforatedbp package itself is missing 34 if e.name == "perforatedbp": 35 pass 36 else: 37 # perforatedbp exists but is missing a dependency 38 raise 39 40import copy 41 42from safetensors.torch import load_file 43from safetensors.torch import save_file 44from safetensors.torch import safe_open 45 46 47def perforate_model( 48 model, 49 doing_pai=True, 50 save_name="", 51 making_graphs=True, 52 maximizing_score=None, 53 num_classes=10000000000, 54 values_per_train_epoch=-1, 55 values_per_val_epoch=-1, 56 zooming_graph=True, 57 config_file=None, 58): 59 """Main function to initialize the network to add dendrites 60 61 This kicks off the entire Perforated AI process to add 62 the scaffolding to the network to be able to add dendrites 63 64 Parameters 65 ---------- 66 model : nn.Module 67 The neural network model to initialize. 68 doing_pai : bool, optional 69 Whether to actually add dendrites, by default True 70 save_name : str, optional 71 The name to save the model under, by default "PAI" 72 making_graphs : bool, optional 73 Whether to create graphs during training, by default True 74 maximizing_score : bool or None, optional 75 Whether to maximize the score during training. 76 - None: use config value (defaults to True) 77 - bool: override config value for this run 78 Setting False is for when the score is a loss to be minimized. 79 num_classes : int, optional 80 The number of output classes, unused in current version 81 values_per_train_epoch : int, optional 82 The number of values to look back for graphing 83 during training, by default -1 (all values). 84 values_per_val_epoch : int, optional 85 The number of values to look back for graphing 86 during validation, by default -1 (all values). 87 zooming_graph : bool, optional 88 Whether to enable zooming on the graphs, by default True 89 config_file : str or None, optional 90 Optional local JSON config file path to load, by default None. 91 92 Returns 93 ------- 94 model : nn.Module 95 The modified model with dendrite scaffolding added if doing_pai is True 96 97 """ 98 99 if save_name == "": 100 if GPA.pc.get_save_name() == "": 101 save_name = "PAI" 102 else: 103 save_name = GPA.pc.get_save_name() 104 105 if "/" in save_name: 106 print( 107 f"Warning: save_name '{save_name}' contains '/'. Relative paths are not implemented yet." 108 ) 109 sys.exit(1) 110 111 sanitized_save_name = "".join( 112 ch for ch in save_name if ch.isalnum() or ch in ("_", "-", ".") 113 ) 114 if sanitized_save_name != save_name: 115 print( 116 f"Warning: save_name '{save_name}' contained spaces or special characters. " 117 f"Using '{sanitized_save_name}' instead." 118 ) 119 save_name = sanitized_save_name 120 121 if save_name == "": 122 print("Warning: save_name became empty after sanitization. Using 'PAI'.") 123 save_name = "PAI" 124 125 if config_file is not None: 126 GPA.pc.set_config_file(config_file) 127 128 GPA.pc.set_save_name(save_name) 129 GPA.pc.sync_config_sources() 130 131 if maximizing_score is None: 132 maximizing_score = GPA.pc.get_maximizing_score() 133 else: 134 # Programmatic perforate_model input must win over config-file values. 135 GPA.pc.set_maximizing_score(maximizing_score) 136 137 if not GPA.pc.get_configuration_confirmed(): 138 CPA.set_perforation_targets(model) 139 140 141 GPA.pai_tracker = TPA.PAINeuronModuleTracker( 142 doing_pai=doing_pai, save_name=save_name 143 ) 144 if _dashboard_emitter is not None: 145 _dashboard_emitter.emit_run_start(GPA.pc, save_name) 146 147 # Get config outputs before tracker is initialized 148 GPA.pc.persist_config_outputs(overwrite_config_file=False) 149 150 model = GPA.pai_tracker.initialize( 151 model, 152 doing_pai=doing_pai, 153 save_name=save_name, 154 making_graphs=making_graphs, 155 maximizing_score=maximizing_score, 156 num_classes=num_classes, 157 values_per_train_epoch=-values_per_train_epoch, 158 values_per_val_epoch=values_per_val_epoch, 159 zooming_graph=zooming_graph, 160 ) 161 162 163 164 return model 165 166 167def get_pai_modules(net, depth, seen_ids=None): 168 """Get a list of all neuron modules 169 170 Parameters 171 ---------- 172 net : nn.Module 173 The module to search. 174 depth : int 175 The current depth in the recursion. 176 177 Returns 178 ------- 179 list 180 A list of all PAI neuron modules found in the network. 181 182 """ 183 if seen_ids is None: 184 seen_ids = set() 185 all_members = net.__dir__() 186 this_list = [] 187 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 188 for submodule_id, layer in net.named_children(): 189 # If there is a self pointer ignore it 190 if net.get_submodule(submodule_id) is net: 191 continue 192 if type(net.get_submodule(submodule_id)) is PA.PAINeuronModule: 193 module = net.get_submodule(submodule_id) 194 if id(module) in seen_ids: 195 continue 196 seen_ids.add(id(module)) 197 this_list = this_list + [module] 198 else: 199 this_list = this_list + get_pai_modules( 200 net.get_submodule(submodule_id), depth + 1, seen_ids 201 ) 202 else: 203 for member in all_members: 204 if isinstance(getattr(type(net), member, None), property): 205 continue 206 # if the getter fails or it is a self pointer ignore it 207 try: 208 if getattr(net, member, None) is net: 209 continue 210 except: 211 continue 212 if type(getattr(net, member, None)) is PA.PAINeuronModule: 213 module = getattr(net, member) 214 if id(module) in seen_ids: 215 continue 216 seen_ids.add(id(module)) 217 this_list = this_list + [module] 218 elif ( 219 issubclass(type(getattr(net, member, None)), nn.Module) 220 or issubclass(type(getattr(net, member, None)), nn.Sequential) 221 or issubclass(type(getattr(net, member, None)), nn.ModuleList) 222 ): 223 this_list = this_list + get_pai_modules( 224 getattr(net, member), depth + 1, seen_ids 225 ) 226 227 return this_list 228 229 230def get_tracked_modules(net, depth, seen_ids=None): 231 """Get a list of all tracked modules 232 233 Parameters 234 ---------- 235 net : nn.Module 236 The module to search. 237 depth : int 238 The current depth in the recursion. 239 240 Returns 241 ------- 242 list 243 A list of all tracked modules found in the network. 244 245 """ 246 if seen_ids is None: 247 seen_ids = set() 248 all_members = net.__dir__() 249 this_list = [] 250 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 251 for submodule_id, layer in net.named_children(): 252 if net.get_submodule(submodule_id) is net: 253 continue 254 if type(net.get_submodule(submodule_id)) is PA.TrackedNeuronModule: 255 module = net.get_submodule(submodule_id) 256 if id(module) in seen_ids: 257 continue 258 seen_ids.add(id(module)) 259 this_list = this_list + [module] 260 else: 261 this_list = this_list + get_tracked_modules( 262 net.get_submodule(submodule_id), depth + 1, seen_ids 263 ) 264 else: 265 for member in all_members: 266 if isinstance(getattr(type(net), member, None), property): 267 continue 268 # if the getter fails or it is a self pointer ignore it 269 try: 270 if getattr(net, member, None) is net: 271 continue 272 except: 273 continue 274 if type(getattr(net, member, None)) is PA.TrackedNeuronModule: 275 module = getattr(net, member) 276 if id(module) in seen_ids: 277 continue 278 seen_ids.add(id(module)) 279 this_list = this_list + [module] 280 elif issubclass(type(getattr(net, member, None)), nn.Module): 281 this_list = this_list + get_tracked_modules( 282 getattr(net, member), depth + 1, seen_ids 283 ) 284 return this_list 285 286 287def get_pai_module_params(net, depth, seen_ids=None): 288 """Get a list of all neuron module parameters 289 290 Parameters 291 ---------- 292 net : nn.Module 293 The module to search. 294 depth : int 295 The current depth in the recursion. 296 297 Returns 298 ------- 299 list 300 A list of all parameters of neuron modules found in this module. 301 302 """ 303 304 if seen_ids is None: 305 seen_ids = set() 306 all_members = net.__dir__() 307 this_list = [] 308 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 309 for submodule_id, layer in net.named_children(): 310 if isinstance(net.get_submodule(submodule_id), PA.PAINeuronModule): # 311 module = net.get_submodule(submodule_id) 312 if id(module) in seen_ids: 313 continue 314 seen_ids.add(id(module)) 315 for param in module.parameters(): 316 if param.requires_grad: 317 this_list = this_list + [param] 318 else: 319 this_list = this_list + get_pai_module_params( 320 net.get_submodule(submodule_id), depth + 1, seen_ids 321 ) 322 else: 323 for member in all_members: 324 if isinstance(getattr(type(net), member, None), property): 325 continue 326 if getattr(net, member, None) == net: 327 continue 328 if isinstance(getattr(net, member, None), PA.PAINeuronModule): 329 module = getattr(net, member) 330 if id(module) in seen_ids: 331 continue 332 seen_ids.add(id(module)) 333 for param in module.parameters(): 334 if param.requires_grad: 335 this_list = this_list + [param] 336 elif issubclass(type(getattr(net, member, None)), nn.Module): 337 this_list = this_list + get_pai_module_params( 338 getattr(net, member), depth + 1, seen_ids 339 ) 340 return this_list 341 342 343def get_pai_network_params(net): 344 """Get a list of all neuron module parameters 345 346 Parameters 347 ---------- 348 net : nn.Module 349 The full model to search. 350 351 Returns 352 ------- 353 list 354 A list of all parameters of neuron modules found in the network. 355 356 """ 357 param_list = get_pai_module_params(net, 0) 358 return param_list 359 360 361def replace_predefined_modules(start_module): 362 """Replace a module with the module from globals list 363 364 Parameters 365 ---------- 366 start_module : nn.Module 367 The module to replace. 368 369 Returns 370 ------- 371 nn.Module 372 The replaced module. 373 374 """ 375 index = GPA.pc.get_modules_to_replace().index(type(start_module)) 376 return GPA.pc.get_replacement_modules()[index](start_module) 377 378 379def scan_module_aliases(net): 380 """Find alias module paths that point to already-seen module instances. 381 382 Parameters 383 ---------- 384 net : Any PyTorch Module. 385 386 Returns 387 ------- 388 dict[str, str] 389 Mapping of alias module paths to their canonical module paths. 390 """ 391 canonical = {} 392 aliases = {} 393 for name, module in net.named_modules(remove_duplicate=False): 394 if name == "": 395 continue 396 sub_name = "." + name 397 module_id = id(module) 398 if module_id in canonical: 399 aliases[sub_name] = canonical[module_id] 400 else: 401 canonical[module_id] = sub_name 402 return aliases 403 404 405def convert_module( 406 net, 407 depth, 408 name_so_far, 409 converted_list, 410 converted_names_list, 411 neuron_module_class, 412 tracked_module_class, 413 config_setup=False, 414): 415 """Recursive function to do all conversion of modules to wrappers of modules 416 417 This is the function that goes through all of the module lists from 418 the globals file and does all the conversion and replacements to 419 setup the dendrite scaffolding as instructed. 420 421 Parameters 422 ---------- 423 net : nn.Module 424 The module to convert. 425 depth : int 426 The current depth in the recursion. 427 name_so_far : str 428 The name of the module so far in the recursion. 429 converted_list : list 430 A list of already converted module ids to avoid infinite loops. 431 converted_names_list : list 432 A corresponding list to help debug duplicate conversions 433 434 Returns 435 ------- 436 nn.Module 437 The converted module. 438 439 """ 440 if GPA.pc.get_verbose(): 441 print("calling convert on %s depth %d" % (net, depth)) 442 print( 443 "calling convert on %s: %s, depth %d" 444 % (name_so_far, type(net).__name__, depth) 445 ) 446 if isinstance(net, neuron_module_class) or ( 447 (tracked_module_class is not None) and isinstance(net, tracked_module_class) 448 ): 449 if GPA.pc.get_verbose(): 450 print( 451 "This is only being called because something in your model " 452 "is pointed to twice by two different variables. Highest " 453 "thing on the list is one of the duplicates" 454 ) 455 return net 456 if depth == 0 and name_so_far == "": 457 aliases = scan_module_aliases(net) 458 existing_not_save = set(GPA.pc.get_module_names_to_not_save()) 459 aliases_to_skip = [ 460 alias for alias in aliases.keys() if alias not in existing_not_save 461 ] 462 if aliases_to_skip and not config_setup: 463 GPA.pc.append_module_names_to_not_save(aliases_to_skip) 464 print( 465 "Auto-detected duplicate module aliases via named_modules; " 466 "keeping first-seen paths and skipping:" 467 ) 468 for alias in aliases_to_skip: 469 print(" - %s (keeps %s)" % (alias, aliases[alias])) 470 all_members = net.__dir__() 471 if GPA.pc.get_extra_verbose(): 472 print("all members:") 473 for member in all_members: 474 print(" - %s" % member) 475 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 476 for submodule_id, layer in net.named_children(): 477 sub_name = name_so_far + "." + str(submodule_id) 478 if sub_name in GPA.pc.get_module_ids_to_track(): 479 if GPA.pc.get_verbose(): 480 print("Seq ID is in track IDs: %s" % sub_name) 481 if tracked_module_class is None: 482 continue 483 if config_setup: 484 net.get_submodule(submodule_id).pai_mode = "tracked" 485 else: 486 setattr( 487 net, 488 submodule_id, 489 tracked_module_class(net.get_submodule(submodule_id), sub_name), 490 ) 491 continue 492 if sub_name in GPA.pc.get_module_ids_to_perforate(): 493 if GPA.pc.get_verbose(): 494 print("Seq ID is in convert IDs: %s" % sub_name) 495 if config_setup: 496 net.get_submodule(submodule_id).pai_mode = "perforated" 497 else: 498 setattr( 499 net, 500 submodule_id, 501 neuron_module_class(net.get_submodule(submodule_id), sub_name), 502 ) 503 continue 504 if type(net.get_submodule(submodule_id)) in GPA.pc.get_modules_to_replace(): 505 if GPA.pc.get_verbose(): 506 print( 507 "Seq sub is in replacement module so replacing: %s" % sub_name 508 ) 509 if config_setup: 510 net.get_submodule(submodule_id).pai_mode = "replaced" 511 else: 512 setattr( 513 net, 514 submodule_id, 515 replace_predefined_modules(net.get_submodule(submodule_id)), 516 ) 517 if ( 518 type(net.get_submodule(submodule_id)) in GPA.pc.get_modules_to_track() 519 ) or ( 520 type(net.get_submodule(submodule_id)).__name__ 521 in GPA.pc.get_module_names_to_track() 522 ): 523 if GPA.pc.get_verbose(): 524 print( 525 "Seq sub is in tracking list so initiating tracked for: %s" 526 % sub_name 527 ) 528 if tracked_module_class is None: 529 continue 530 if config_setup: 531 net.get_submodule(submodule_id).pai_mode = "tracked" 532 else: 533 setattr( 534 net, 535 submodule_id, 536 tracked_module_class(net.get_submodule(submodule_id), sub_name), 537 ) 538 elif ( 539 type(net.get_submodule(submodule_id)) 540 in GPA.pc.get_modules_to_perforate() 541 or type(net.get_submodule(submodule_id)).__name__ 542 in GPA.pc.get_module_names_to_perforate() 543 ): 544 if GPA.pc.get_verbose(): 545 print( 546 "Seq sub is in conversion list so initing PAI for: " 547 "%s" % sub_name 548 ) 549 if ( 550 issubclass( 551 type(net.get_submodule(submodule_id)), 552 torch.nn.modules.batchnorm._BatchNorm, 553 ) 554 or issubclass( 555 type(net.get_submodule(submodule_id)), 556 torch.nn.modules.instancenorm._InstanceNorm, 557 ) 558 or issubclass( 559 type(net.get_submodule(submodule_id)), 560 torch.nn.modules.normalization.LayerNorm, 561 ) 562 ): 563 print( 564 "You have an unwrapped normalization layer, this " 565 "is not recommended: " + name_so_far 566 ) 567 pdb.set_trace() 568 if config_setup: 569 net.get_submodule(submodule_id).pai_mode = "perforated" 570 else: 571 setattr( 572 net, 573 submodule_id, 574 neuron_module_class(net.get_submodule(submodule_id), sub_name), 575 ) 576 else: 577 if net != net.get_submodule(submodule_id): 578 converted_list += [id(net.get_submodule(submodule_id))] 579 converted_names_list += [sub_name] 580 if GPA.pc.get_verbose(): 581 print( 582 "sub is module but in no lists so going deeper: %s" 583 % sub_name 584 ) 585 586 setattr( 587 net, 588 submodule_id, 589 convert_module( 590 net.get_submodule(submodule_id), 591 depth + 1, 592 sub_name, 593 converted_list, 594 converted_names_list, 595 neuron_module_class, 596 tracked_module_class, 597 config_setup, 598 ), 599 ) 600 # else: 601 # print('%s is a self pointer so skipping' % (name_so_far + '[' + str(submodule_id) + ']')) 602 elif type(net) in GPA.pc.get_modules_to_track(): 603 # print('skipping type for returning from call to: %s' % (name_so_far)) 604 return net 605 else: 606 for member in all_members: 607 if isinstance(getattr(type(net), member, None), property): 608 continue 609 # Immediately check if able to get the member, if not skip it 610 try: 611 getattr(net, member, None) 612 except: 613 continue 614 sub_name = name_so_far + "." + member 615 member_obj = getattr(net, member, None) 616 617 if isinstance(member_obj, (torch.nn.Parameter, torch.nn.parameter.Parameter)): 618 if sub_name in GPA.pc.get_parameter_ids_to_track(): 619 if GPA.pc.get_verbose(): 620 print("tracking parameter by ID: %s" % sub_name) 621 if config_setup: 622 member_obj.pai_mode = "tracked" 623 else: 624 member_obj.parameter_type = "neuron" 625 member_obj.wrapped = True 626 continue 627 628 # Track module object ids once at this level so duplicate aliases are 629 # caught consistently (including direct children of the root module). 630 if isinstance(member_obj, nn.Module): 631 if id(member_obj) in converted_list: 632 original_sub_name = converted_names_list[ 633 converted_list.index(id(member_obj)) 634 ] 635 print( 636 "The following module has a duplicate pointer within " 637 "your model: %s" % sub_name 638 ) 639 print("Keeping first pointer: %s" % original_sub_name) 640 print("Skipping duplicate pointer: %s" % sub_name) 641 print( 642 "If you prefer to keep %s and skip %s, add %s to module_names_to_not_save before convert." 643 % (sub_name, original_sub_name, original_sub_name) 644 ) 645 if not config_setup: 646 GPA.pc.append_module_names_to_not_save([sub_name]) 647 continue 648 converted_list += [id(member_obj)] 649 converted_names_list += [sub_name] 650 if sub_name in GPA.pc.get_module_ids_to_track(): 651 if GPA.pc.get_verbose(): 652 print("Seq ID is in track IDs: %s" % sub_name) 653 if tracked_module_class is None: 654 continue 655 if config_setup: 656 getattr(net, member).pai_mode = "tracked" 657 else: 658 setattr( 659 net, member, tracked_module_class(getattr(net, member), sub_name) 660 ) 661 continue 662 if sub_name in GPA.pc.get_module_ids_to_perforate(): 663 if GPA.pc.get_verbose(): 664 print("Seq ID is in convert IDs: %s" % sub_name) 665 if config_setup: 666 getattr(net, member).pai_mode = "perforated" 667 else: 668 setattr( 669 net, member, neuron_module_class(getattr(net, member), sub_name) 670 ) 671 continue 672 if id(getattr(net, member, None)) == id(net): 673 if GPA.pc.get_verbose(): 674 print("member sub is a self pointer: %s" % sub_name) 675 continue 676 if sub_name in GPA.pc.get_module_names_to_not_save(): 677 if GPA.pc.get_verbose(): 678 print("Skipping %s during convert" % sub_name) 679 else: 680 if sub_name == ".base_model": 681 print( 682 "By default skipping base_model. See " 683 '"Safetensors Errors" section of ' 684 "customization.md to include it." 685 ) 686 continue 687 if type(getattr(net, member, None)) in GPA.pc.get_modules_to_replace(): 688 if GPA.pc.get_verbose(): 689 print("sub is in replacement module so replacing: %s" % sub_name) 690 if config_setup: 691 getattr(net, member).pai_mode = "replaced" 692 else: 693 setattr( 694 net, member, replace_predefined_modules(getattr(net, member, None)) 695 ) 696 if ( 697 type(getattr(net, member, None)) in GPA.pc.get_modules_to_track() 698 or type(getattr(net, member, None)).__name__ 699 in GPA.pc.get_module_names_to_track() 700 or sub_name in GPA.pc.get_module_ids_to_track() 701 ): 702 if GPA.pc.get_verbose(): 703 print( 704 "sub is in tracking list so initiating tracked for: %s" 705 % sub_name 706 ) 707 if tracked_module_class is None: 708 continue 709 if config_setup: 710 getattr(net, member).pai_mode = "tracked" 711 else: 712 setattr( 713 net, member, tracked_module_class(getattr(net, member), sub_name) 714 ) 715 elif ( 716 type(getattr(net, member, None)) in GPA.pc.get_modules_to_perforate() 717 or type(getattr(net, member, None)).__name__ 718 in GPA.pc.get_module_names_to_perforate() 719 or (sub_name in GPA.pc.get_module_ids_to_perforate()) 720 ): 721 if GPA.pc.get_verbose(): 722 print( 723 "sub is in conversion list so initiating PAI for: %s" % sub_name 724 ) 725 if config_setup: 726 getattr(net, member).pai_mode = "perforated" 727 else: 728 setattr( 729 net, 730 member, 731 neuron_module_class(getattr(net, member), sub_name), 732 ) 733 elif ( 734 issubclass(type(getattr(net, member, None)), nn.Module) 735 or issubclass(type(getattr(net, member, None)), nn.Sequential) 736 or issubclass(type(getattr(net, member, None)), nn.ModuleList) 737 ): 738 if net != getattr(net, member): 739 if GPA.pc.get_verbose(): 740 print( 741 "sub is module but in no lists so going deeper: %s" 742 % sub_name 743 ) 744 setattr( 745 net, 746 member, 747 convert_module( 748 getattr(net, member), 749 depth + 1, 750 sub_name, 751 converted_list, 752 converted_names_list, 753 neuron_module_class, 754 tracked_module_class, 755 config_setup, 756 ), 757 ) 758 if ( 759 issubclass( 760 type(getattr(net, member, None)), 761 torch.nn.modules.batchnorm._BatchNorm, 762 ) 763 or issubclass( 764 type(getattr(net, member, None)), 765 torch.nn.modules.instancenorm._InstanceNorm, 766 ) 767 or issubclass( 768 type(getattr(net, member, None)), 769 torch.nn.modules.normalization.LayerNorm, 770 ) 771 ): 772 if not GPA.pc.get_unwrapped_modules_confirmed(): 773 print( 774 "potentially found a norm Layer that " 775 "is not accounted for, this is not recommended: %s" % (sub_name) 776 ) 777 print( 778 "Set GPA.pc.set_unwrapped_modules_confirmed(True) to skip " 779 "this next time only if you are unable to debug" 780 ) 781 print( 782 "inspect your network to " 783 "see what the module type containing this layer is." 784 ) 785 print("Then do one of the following:") 786 print( 787 " - Add the module type to " 788 "GPA.pc.get_module_names_to_perforate() to wrap it entirely" 789 ) 790 print( 791 " - If the norm layer is part of a sequential wrap " 792 "it and the previous layer in a PAISequential" 793 ) 794 print( 795 " - If you do not want to add dendrites to this " 796 "module add the type to GPA.pc.get_module_names_to_track()" 797 ) 798 pdb.set_trace() 799 else: 800 if GPA.pc.get_verbose(): 801 if member[0] != "_" or GPA.pc.get_extra_verbose() is True: 802 print("not calling convert on %s depth %d" % (member, depth)) 803 if GPA.pc.get_verbose(): 804 print("returning from call to: %s" % (name_so_far)) 805 return net 806 807 808def convert_network(net, layer_name=""): 809 """Function that calls convert_module and checks results 810 811 Parameters 812 ---------- 813 net : nn.Module 814 The network to convert. 815 layer_name : str, optional 816 The name of the layer if converting a single layer, by default "" 817 818 Returns 819 ------- 820 nn.Module 821 The converted network. 822 823 """ 824 if GPA.pc.get_perforated_backpropagation(): 825 UPB.initialize_pb() 826 MPB.set_main_parameters(net) 827 if type(net) in GPA.pc.get_modules_to_replace(): 828 net = replace_predefined_modules(net) 829 if (type(net) in GPA.pc.get_modules_to_perforate()) or ( 830 type(net).__name__ in GPA.pc.get_module_names_to_perforate() 831 ): 832 if layer_name == "": 833 print( 834 "converting a single layer without a name, add a " 835 "layer_name param to the call" 836 ) 837 sys.exit(-1) 838 net = PA.PAINeuronModule(net, layer_name) 839 else: 840 net = convert_module( 841 net, 0, "", [], [], PA.PAINeuronModule, PA.TrackedNeuronModule 842 ) 843 if GPA.pai_tracker.member_vars["doing_pai"]: 844 missed_ones = [] 845 tracked_ones = [] 846 for name, param in net.named_parameters(): 847 wrapped = "wrapped" in param.__dir__() 848 if wrapped: 849 if GPA.pc.get_verbose(): 850 print("param %s is now wrapped" % (name)) 851 else: 852 tracked = "tracked" in param.__dir__() 853 if tracked: 854 tracked_ones.append(name) 855 else: 856 missed_ones.append(name) 857 if ( 858 len(missed_ones) != 0 or len(tracked_ones) != 0 859 ) and GPA.pc.get_unwrapped_modules_confirmed() is False: 860 print( 861 "\n------------------------------------------------------------------" 862 ) 863 print( 864 "The following params are not wrapped.\n------------------------------------------------------------------" 865 ) 866 for name in tracked_ones: 867 print("." + name) 868 print( 869 "\n------------------------------------------------------------------" 870 ) 871 print( 872 "The following params are not tracked or wrapped.\n------------------------------------------------------------------" 873 ) 874 for name in missed_ones: 875 print("." + name) 876 print( 877 "\n------------------------------------------------------------------" 878 ) 879 print( 880 "Modules that are not wrapped will not have Dendrites to optimize them" 881 ) 882 print( 883 "Modules modules that are not tracked can cause errors and is NOT recommended" 884 ) 885 print( 886 "Any modules in the second list should be added to module_names_to_track" 887 ) 888 889 print( 890 "Set GPA.pc.set_unwrapped_modules_confirmed(True) to skip this next time" 891 ) 892 print( 893 "Inspect your network and see what the module types of these values are to add them to PGB.module_names_to_perforate" 894 ) 895 # If did miss some then set trace to debug 896 if len(missed_ones) != 0: 897 print( 898 "------------------------------------------------------------------\nType 'c' + enter to continue the run to confirm you do not want them to be refined" 899 ) 900 901 pdb.set_trace() 902 print("confirmed") 903 net.register_buffer("tracker_string", torch.tensor([], dtype=torch.uint8)) 904 return net 905 906 907def string_to_tensor(string): 908 """Helper function to convert a layer_tracker into a string 909 910 This is required for safetensors saving 911 912 Parameters 913 ---------- 914 string : str 915 The string to convert. 916 917 Returns 918 ------- 919 torch.Tensor 920 The converted tensor. 921 922 """ 923 ords = list(map(ord, string)) 924 ords = torch.tensor(ords, dtype=torch.uint8) 925 return ords 926 927 928def string_from_tensor(string_tensor): 929 """Convert a tensor back into a string 930 931 Parameters 932 ---------- 933 string_tensor : torch.Tensor 934 The tensor to convert. 935 936 Returns 937 ------- 938 str 939 The converted string. 940 941 """ 942 ords = string_tensor.tolist() 943 to_return = "" 944 # Doing block processing like this helps with memory errors 945 while len(ords) != 0: 946 remaining_ords = ords[100000:] 947 ords = ords[:100000] 948 to_append = "".join(map(chr, ords)) 949 to_return = to_return + to_append 950 ords = remaining_ords 951 return to_return 952 953 954def save_system(net, folder, name): 955 """Save the entire system 956 957 This saves the network itself as well as the tracker information 958 959 Parameters 960 ---------- 961 net : nn.Module 962 The network to save. 963 folder : str 964 The folder to save the network in. 965 name : str 966 The name to save the network under. 967 968 Returns 969 ------- 970 None 971 972 """ 973 if GPA.pc.get_verbose(): 974 print("saving system %s" % name) 975 temp = string_to_tensor(GPA.pai_tracker.to_string()) 976 if hasattr(net, "tracker_string"): 977 net.tracker_string = string_to_tensor(GPA.pai_tracker.to_string()).to( 978 next(net.parameters()).device 979 ) 980 else: 981 net.register_buffer( 982 "tracker_string", 983 string_to_tensor(GPA.pai_tracker.to_string()).to( 984 next(net.parameters()).device 985 ), 986 ) 987 # Before saving the tracker must be cleared to not contain pointers to the 988 # models modules 989 old_list = GPA.pai_tracker.neuron_module_vector 990 GPA.pai_tracker.neuron_module_vector = [] 991 save_net(net, folder, name) 992 GPA.pai_tracker.neuron_module_vector = old_list 993 pai_save_system(net, folder, name) 994 995 996def load_system( 997 net, 998 folder, 999 name, 1000 load_from_restart=False, 1001 switch_call=False, 1002 load_from_manual_save=False, 1003): 1004 """Load the entire system 1005 1006 This is what should be used to load a saved system and restart training 1007 1008 Parameters 1009 ---------- 1010 net : nn.Module 1011 The network to load into. 1012 folder : str 1013 The folder to load the network from. 1014 name : str 1015 The name to load the network from. 1016 load_from_restart : bool, optional 1017 Whether this is being loaded from an automatic restart, by default False 1018 switch_call : bool, optional 1019 Whether this is being called from a switch, by default False 1020 load_from_manual_save : bool, optional 1021 Whether this is being loaded from a manual save, by default False 1022 1023 Returns 1024 ------- 1025 nn.Module 1026 The loaded network. 1027 1028 Notes 1029 ----- 1030 If you manually call save_system then load_from_manual_save should be True 1031 1032 """ 1033 if GPA.pc.get_verbose(): 1034 print("loading system %s" % name) 1035 net = load_net(net, folder, name) 1036 GPA.pai_tracker.reset_module_vector(net, load_from_restart) 1037 1038 GPA.pai_tracker.from_string(string_from_tensor(net.tracker_string)) 1039 GPA.pai_tracker.saved_time = time.time() 1040 GPA.pai_tracker.loaded = True 1041 GPA.pai_tracker.member_vars["current_best_validation_score"] = 0 1042 GPA.pai_tracker.member_vars["epoch_last_improved"] = GPA.pai_tracker.member_vars[ 1043 "num_epochs_run" 1044 ] 1045 if GPA.pc.get_verbose(): 1046 print( 1047 "after loading epoch last improved is %d mode is %c" 1048 % ( 1049 GPA.pai_tracker.member_vars["epoch_last_improved"], 1050 GPA.pai_tracker.member_vars["mode"], 1051 ) 1052 ) 1053 1054 # Saves always take place before the call to start_epoch so call it here 1055 # when loading to correct off by 1 problems 1056 if (not switch_call) and (not load_from_manual_save): 1057 GPA.pai_tracker.start_epoch(internal_call=True) 1058 return net 1059 1060 1061def load_pretrained_model( 1062 net, 1063 folder, 1064 name, 1065 remove_dendrite_scaffolding=False, 1066): 1067 """Load a pretrained perforated model and reset tracker for fresh training. 1068 1069 This function loads a pretrained model's weights and dendrite structure while 1070 resetting all tracker state (epochs, switch history, etc.) to start training 1071 from scratch on a new task. This is useful for transfer learning where you want 1072 pretrained weights but need fresh training dynamics. 1073 1074 Parameters 1075 ---------- 1076 net : nn.Module 1077 The network to load into. 1078 folder : str 1079 The folder containing the pretrained model. 1080 name : str 1081 The name of the checkpoint to load (e.g., 'best_model', 'beforeSwitch_0'). 1082 remove_dendrite_scaffolding : bool, optional 1083 If True, removes dendrite scaffolding for inference or finetuning without 1084 adding more dendrites using blockwise_network and refresh_net. Default False. 1085 1086 Returns 1087 ------- 1088 nn.Module 1089 The loaded network with reset tracker state. 1090 1091 Examples 1092 -------- 1093 Load pretrained weights for continued dendrite training: 1094 >>> model = load_pretrained_model(model, "pretrained-prefc", "beforeSwitch_0") 1095 1096 Load pretrained weights for finetuning without adding more dendrites: 1097 >>> model = load_pretrained_model(model, "pretrained-prefc", "best_model", 1098 ... remove_dendrite_scaffolding=True) 1099 1100 Notes 1101 ----- 1102 This function: 1103 - Loads model weights and dendrite structure from checkpoint 1104 - Resets all epoch counters to -1 (will become 0 after first start_epoch) 1105 - Resets switch history and validation score tracking 1106 - Clears accuracy/loss history arrays 1107 - Optionally removes dendrite scaffolding (no more dendrite additions) 1108 1109 The tracker is reset to behave as if starting fresh training, while keeping 1110 the learned weights and dendrite structure from the pretrained model. 1111 """ 1112 from perforatedai import globals_perforatedai as GPA 1113 1114 if GPA.pc.get_verbose(): 1115 print(f"Loading pretrained model from {folder}/{name}") 1116 1117 # Load the model weights and dendrite structure 1118 net = load_system(net, folder, name, load_from_manual_save=True) 1119 1120 if GPA.pc.get_verbose(): 1121 print("Resetting tracker state for fresh training...") 1122 1123 # Reset structural training state to true initial values. 1124 # Keeping pretrained architecture/weights while zeroing cycle counters avoids 1125 # stale dendrite bookkeeping referencing empty score buffers. 1126 GPA.pai_tracker.reset_module_vector(net, load_from_restart=True) 1127 GPA.pai_tracker.member_vars["mode"] = "n" 1128 GPA.pai_tracker.member_vars["num_dendrites_added"] = 0 1129 GPA.pai_tracker.member_vars["num_dendrites_integrated"] = 0 1130 GPA.pai_tracker.member_vars["num_cycles"] = 0 1131 GPA.pai_tracker.member_vars["num_dendrite_tries"] = 0 1132 GPA.pai_tracker.member_vars["current_n_set_global_best"] = True 1133 1134 # Reset epoch counters 1135 GPA.pai_tracker.member_vars["num_epochs_run"] = -1 1136 GPA.pai_tracker.member_vars["total_epochs_run"] = -1 1137 GPA.pai_tracker.member_vars["epoch_last_improved"] = 0 1138 GPA.pai_tracker.member_vars["last_switch"] = 0 1139 GPA.pai_tracker.member_vars["manual_train_switch"] = False 1140 1141 # Reset switch history 1142 GPA.pai_tracker.member_vars["switch_epochs"] = [] 1143 GPA.pai_tracker.member_vars["n_switch_epochs"] = [] 1144 GPA.pai_tracker.member_vars["p_switch_epochs"] = [] 1145 GPA.pai_tracker.member_vars["param_counts"] = [] 1146 1147 # Reset validation scores and tracking 1148 GPA.pai_tracker.member_vars["current_best_validation_score"] = 0 1149 GPA.pai_tracker.member_vars["global_best_validation_score"] = 0 1150 GPA.pai_tracker.member_vars["running_accuracy"] = 0 1151 1152 # Clear accuracy/loss history arrays 1153 GPA.pai_tracker.member_vars["accuracies"] = [] 1154 GPA.pai_tracker.member_vars["last_improved_accuracies"] = [] 1155 GPA.pai_tracker.member_vars["test_accuracies"] = [] 1156 GPA.pai_tracker.member_vars["n_accuracies"] = [] 1157 GPA.pai_tracker.member_vars["p_accuracies"] = [] 1158 GPA.pai_tracker.member_vars["running_accuracies"] = [] 1159 GPA.pai_tracker.member_vars["training_loss"] = [] 1160 GPA.pai_tracker.member_vars["training_learning_rates"] = [] 1161 GPA.pai_tracker.member_vars["test_scores"] = [] 1162 1163 # Clear extra scores 1164 GPA.pai_tracker.member_vars["extra_scores"] = {} 1165 GPA.pai_tracker.member_vars["extra_scores_without_graphing"] = {} 1166 GPA.pai_tracker.member_vars["n_extra_scores"] = {} 1167 1168 # Keep per-layer dendrite score buffers initialized by reset_module_vector. 1169 1170 # Clear timing arrays 1171 GPA.pai_tracker.member_vars["n_epoch_times"] = [] 1172 GPA.pai_tracker.member_vars["p_epoch_times"] = [] 1173 GPA.pai_tracker.member_vars["n_train_times"] = [] 1174 GPA.pai_tracker.member_vars["p_train_times"] = [] 1175 GPA.pai_tracker.member_vars["n_val_times"] = [] 1176 GPA.pai_tracker.member_vars["p_val_times"] = [] 1177 1178 # Clear overwritten tracking 1179 GPA.pai_tracker.member_vars["overwritten_extras"] = [] 1180 GPA.pai_tracker.member_vars["overwritten_vals"] = [] 1181 GPA.pai_tracker.member_vars["overwritten_epochs"] = 0 1182 1183 # Reset learning rate search state 1184 GPA.pai_tracker.member_vars["initial_lr_test_epoch_count"] = -1 1185 GPA.pai_tracker.member_vars["current_n_learning_rate_initial_skip_steps"] = 0 1186 GPA.pai_tracker.member_vars["last_max_learning_rate_steps"] = 0 1187 GPA.pai_tracker.member_vars["last_max_learning_rate_value"] = -1 1188 GPA.pai_tracker.member_vars["current_cycle_lr_max_scores"] = [] 1189 GPA.pai_tracker.member_vars["current_step_count"] = 0 1190 GPA.pai_tracker.member_vars["committed_to_initial_rate"] = True 1191 GPA.pai_tracker.member_vars["best_mean_score_improved_this_epoch"] = 0 1192 GPA.pai_tracker.member_vars["step_status"] = TPA.STEP_CLEARED 1193 1194 # Reset saved time 1195 GPA.pai_tracker.start_time = time.time() 1196 GPA.pai_tracker.saved_time = 0 1197 1198 # Match tracker initialization behavior so first validation uses epoch 0. 1199 GPA.pai_tracker.start_epoch(internal_call=True) 1200 1201 if GPA.pc.get_verbose(): 1202 print( 1203 f"Tracker reset complete. Dendrites: {GPA.pai_tracker.member_vars['num_dendrites_integrated']}, " 1204 f"Mode: {GPA.pai_tracker.member_vars['mode']}" 1205 ) 1206 1207 # Optionally remove dendrite scaffolding 1208 if remove_dendrite_scaffolding: 1209 if GPA.pc.get_verbose(): 1210 print("Removing dendrite scaffolding (no dendrite additions)...") 1211 1212 from perforatedai import blockwise_perforatedai as BPA 1213 from perforatedai import clean_perforatedai as CPA 1214 1215 net = BPA.blockwise_network(net) 1216 net = CPA.refresh_net(net) 1217 1218 if GPA.pc.get_verbose(): 1219 print("Dendrite scaffolding removed. Model ready for inference or finetuning.") 1220 1221 return net 1222 1223 1224import json 1225from collections import defaultdict 1226from safetensors.torch import save_file, safe_open 1227import torch 1228 1229 1230def save_model_with_weight_tying(model, filepath): 1231 """Save model with safetensors while handling weight tying automatically 1232 1233 Parameters 1234 ---------- 1235 model : Any PyTorch Module. 1236 filepath : Path to filename. 1237 1238 Returns 1239 ------- 1240 dict[str, str] 1241 Mapping from secondary parameter keys to their primary tied key. 1242 """ 1243 state_dict = model.state_dict() 1244 1245 # Find all weight tied parameters 1246 tensor_to_keys = defaultdict(list) 1247 for key, tensor in state_dict.items(): 1248 # Use tensor data pointer as unique identifier 1249 tensor_id = tensor.data_ptr() 1250 tensor_to_keys[tensor_id].append(key) 1251 1252 # Find tied weights (tensors referenced by multiple keys) 1253 tied_weights = {} 1254 keys_to_remove = set() 1255 for tensor_id, keys in tensor_to_keys.items(): 1256 if len(keys) > 1 and not tensor_id == 0: 1257 # Multiple keys reference the same tensor - this is weight tying 1258 # Sort keys for deterministic ordering 1259 keys = sorted(keys) 1260 primary_key = keys[0] # Keep the first key 1261 for secondary_key in keys[1:]: 1262 tied_weights[secondary_key] = primary_key 1263 keys_to_remove.add(secondary_key) 1264 1265 # Remove tied weights from state_dict (keep only primary references) 1266 filtered_state_dict = { 1267 k: v for k, v in state_dict.items() if k not in keys_to_remove 1268 } 1269 1270 # Create metadata for weight tying information 1271 metadata = {} 1272 if tied_weights: 1273 # Store weight tying info as JSON string in metadata 1274 metadata["weight_tying"] = json.dumps(tied_weights) 1275 save_file(filtered_state_dict, filepath, metadata=metadata) 1276 print(f"Saved model with {len(tied_weights)} weight tying relationships") 1277 return tied_weights 1278 1279 1280def load_model_with_weight_tying(model, filepath): 1281 """Load model from safetensors while restoring weight tying 1282 1283 Parameters 1284 ---------- 1285 model : Any PyTorch module. 1286 filepath : Path to model file. 1287 1288 Returns 1289 ------- 1290 Loaded model. 1291 """ 1292 with safe_open(filepath, framework="pt") as f: 1293 metadata = f.metadata() 1294 state_dict = {key: f.get_tensor(key) for key in f.keys()} 1295 1296 # Restore weight tying if metadata exists 1297 tied_weights = {} 1298 if metadata and "weight_tying" in metadata: 1299 tied_weights = json.loads(metadata["weight_tying"]) 1300 for secondary_key, primary_key in tied_weights.items(): 1301 if primary_key in state_dict: 1302 # Restore the tied reference 1303 state_dict[secondary_key] = state_dict[primary_key] 1304 print(f"Restored weight tying: {secondary_key} -> {primary_key}") 1305 1306 # Handle tracker_string loading with flexible key matching 1307 tracker_key = None 1308 if "tracker_string" in state_dict: 1309 tracker_key = "tracker_string" 1310 else: 1311 # Search for keys containing "tracker_string" 1312 tracker_keys = [key for key in state_dict.keys() if "tracker_string" in key] 1313 if len(tracker_keys) == 1: 1314 tracker_key = tracker_keys[0] 1315 elif len(tracker_keys) > 1: 1316 print(f"Error: Multiple tracker_string keys found: {tracker_keys}") 1317 pdb.set_trace() 1318 else: 1319 print("Error: No tracker_string found in state_dict") 1320 1321 if tracker_key is not None and hasattr(model, "tracker_string"): 1322 model.tracker_string = state_dict[tracker_key] 1323 1324 model.load_state_dict(state_dict) 1325 return model 1326 1327 1328def save_net(net, folder, name): 1329 """Save the network 1330 1331 This is called within save_system after the tracker has been 1332 turned into a single tensor to be saved as a part of the network 1333 1334 Parameters 1335 ---------- 1336 net : nn.Module 1337 The network to save. 1338 folder : str 1339 The folder to save the network in. 1340 name : str 1341 The name to save the network under. 1342 1343 Returns 1344 ------- 1345 None 1346 1347 """ 1348 # If running a DDP only save with first thread 1349 if "RANK" in os.environ: 1350 if int(os.environ["RANK"]) != 0: 1351 return 1352 if not os.path.isdir(folder): 1353 os.makedirs(folder) 1354 save_point = folder + "/" 1355 if not os.path.isdir(save_point): 1356 os.mkdir(save_point) 1357 for param in net.parameters(): 1358 param.data = param.data.contiguous() 1359 if GPA.pc.get_using_safe_tensors(): 1360 if GPA.pc.get_weight_tying_experimental(): 1361 save_model_with_weight_tying(net, save_point + name + ".pt") 1362 else: 1363 # Strip the . so that the naming is the same for everywhere but it works with state_dict naming 1364 not_save = [ns.lstrip('.') for ns in GPA.pc.get_module_names_to_not_save()] 1365 state_dict = {k: v for k, v in net.state_dict().items() 1366 if not any(k.startswith(ns) for ns in not_save)} 1367 save_file(state_dict, save_point + name + ".pt") 1368 else: 1369 torch.save(net, save_point + name + ".pt") 1370 1371 1372def save_pai_net(net, folder, name): 1373 """Save the final pai network 1374 1375 This can be called after training to save the final network 1376 with all scaffolding removed so only the refined weights remain 1377 1378 Parameters 1379 ---------- 1380 net : nn.Module 1381 The network to save. 1382 folder : str 1383 The folder to save the network in. 1384 name : str 1385 The name to save the network under. 1386 1387 Returns 1388 ------- 1389 None 1390 1391 """ 1392 # if running a DDP only save with first thread 1393 if "RANK" in os.environ: 1394 if int(os.environ["RANK"]) != 0: 1395 return 1396 1397 # print('calling save: %s' % name) 1398 # GPA.pai_tracker.archive_layer() 1399 # These deep copys are required or the real model will also have its layers replaced 1400 net = prepare_final_model(net) 1401 if not os.path.isdir(folder): 1402 os.makedirs(folder) 1403 save_point = folder + "/" 1404 if not os.path.isdir(save_point): 1405 os.mkdir(save_point) 1406 1407 if GPA.pc.get_using_safe_tensors(): 1408 if GPA.pc.get_weight_tying_experimental(): 1409 save_model_with_weight_tying(net, save_point + name + "_pai.pt") 1410 else: 1411 save_file(net.state_dict(), save_point + name + "_pai.pt") 1412 else: 1413 torch.save(net, save_point + name + "_pai.pt") 1414 1415 1416def save_pai_net(net, folder, name): 1417 """Save the final pai network 1418 1419 This can be called after training to save the final network 1420 with all scaffolding removed so only the refined weights remain 1421 1422 Parameters 1423 ---------- 1424 net : nn.Module 1425 The network to save. 1426 folder : str 1427 The folder to save the network in. 1428 name : str 1429 The name to save the network under. 1430 1431 Returns 1432 ------- 1433 None 1434 1435 """ 1436 # if running a DDP only save with first thread 1437 if "RANK" in os.environ: 1438 if int(os.environ["RANK"]) != 0: 1439 return 1440 1441 # print('calling save: %s' % name) 1442 # GPA.pai_tracker.archive_layer() 1443 # These deep copys are required or the real model will also have its layers replaced 1444 net = prepare_final_model(net) 1445 if not os.path.isdir(folder): 1446 os.makedirs(folder) 1447 save_point = folder + "/" 1448 if not os.path.isdir(save_point): 1449 os.mkdir(save_point) 1450 1451 if GPA.pc.get_using_safe_tensors(): 1452 if GPA.pc.get_weight_tying_experimental(): 1453 save_model_with_weight_tying(net, save_point + name + "_pai.pt") 1454 else: 1455 save_file(net.state_dict(), save_point + name + "_pai.pt") 1456 else: 1457 torch.save(net, save_point + name + "_pai.pt") 1458 1459 1460def manual_load_state_dict(model, state_dict): 1461 """Load a state dict into a model key-by-key with relaxed checks. 1462 1463 Parameters 1464 ---------- 1465 model : nn.Module 1466 Model receiving parameters. 1467 state_dict : dict 1468 Source parameter dictionary. 1469 1470 Notes 1471 ----- 1472 Keys configured in ``module_names_to_not_save`` are skipped. 1473 1474 Returns 1475 ------- 1476 None 1477 This function does not return a value. 1478 """ 1479 own_state = model.state_dict() 1480 not_save = [ns.lstrip('.') for ns in GPA.pc.get_module_names_to_not_save()] 1481 for name, param in state_dict.items(): 1482 if any(name.startswith(ns) for ns in not_save): 1483 print("skipping loading %s based on module_names_to_not_save" % name) 1484 continue 1485 if name not in own_state: 1486 print(f"Warning: {name} not found in model state_dict") 1487 continue 1488 if isinstance(param, torch.nn.Parameter): 1489 # Backwards compatibility for serialized parameters 1490 param = param.data 1491 try: 1492 own_state[name].copy_(param) 1493 except Exception as e: 1494 print(f"Error loading {name}: {e}") 1495 print("Manual load complete") 1496 1497 1498def load_net(net, folder, name): 1499 """load the network 1500 1501 This is called within load_system after the tracker has been 1502 loaded 1503 1504 Parameters 1505 ---------- 1506 net : nn.Module 1507 The network to save. 1508 folder : str 1509 The folder to save the network in. 1510 name : str 1511 The name to save the network under. 1512 1513 Returns 1514 ------- 1515 nn.Module 1516 The loaded network. 1517 1518 """ 1519 save_point = folder + "/" 1520 if GPA.pc.get_using_safe_tensors(): 1521 model_path = save_point + name + ".pt" 1522 if GPA.pc.get_weight_tying_experimental(): 1523 return load_model_with_weight_tying(net, model_path) 1524 else: 1525 try: 1526 with safe_open(model_path, framework="pt") as f: 1527 metadata = f.metadata() 1528 if metadata and "weight_tying" in metadata: 1529 return load_model_with_weight_tying(net, model_path) 1530 except Exception: 1531 pass 1532 state_dict = load_file(model_path) 1533 else: 1534 # Try three different torch.load call signatures for compatibility across 1535 # different PyTorch versions. All three are attempted before giving up so 1536 # the user can see every error at once and decide how to fix the file. 1537 error_weights_only_false = None 1538 error_no_weights_only = None 1539 error_no_state_dict = None 1540 1541 try: 1542 state_dict = torch.load( 1543 save_point + name + ".pt", 1544 map_location=torch.device("cpu"), 1545 weights_only=False, 1546 ).state_dict() 1547 except Exception as caught_error: 1548 error_weights_only_false = caught_error 1549 1550 if error_weights_only_false is not None: 1551 try: 1552 state_dict = torch.load( 1553 save_point + name + ".pt", 1554 map_location=torch.device("cpu"), 1555 ).state_dict() 1556 except Exception as caught_error: 1557 error_no_weights_only = caught_error 1558 1559 if error_no_weights_only is not None: 1560 try: 1561 state_dict = torch.load( 1562 save_point + name + ".pt", 1563 map_location=torch.device("cpu"), 1564 ) 1565 except Exception as caught_error: 1566 error_no_state_dict = caught_error 1567 1568 if error_no_state_dict is not None: 1569 separator = "\n" + "=" * 60 + "\n" 1570 print(separator.join([ 1571 "\nAll four load attempts failed for: " + save_point + name + ".pt", 1572 "Attempt 1 (weights_only=False, .state_dict()):\n" + str(error_weights_only_false), 1573 "Attempt 2 (default weights_only, .state_dict()):\n" + str(error_no_weights_only), 1574 "Attempt 3 (default weights_only, no .state_dict()):\n" + str(error_no_state_dict), 1575 "You must find a way to fix at least one of the above errors to load successfully.\n", 1576 "Suggested problems to look into:\n", 1577 "1 - Modules which have member variables that are methods, which cant be pickled. See __setstate__ and __getstate__ functions to append." 1578 ])) 1579 raise error_no_state_dict 1580 return load_net_from_dict(net, state_dict) 1581 1582def get_module_base_name(module): 1583 """Normalize a wrapped module name for state-dict key lookup. 1584 1585 Parameters 1586 ---------- 1587 module : nn.Module 1588 Module containing a ``name`` attribute. 1589 1590 Returns 1591 ------- 1592 str 1593 Base name with leading dot and optional ``module.`` prefix removed. 1594 """ 1595 module_name = module.name 1596 # This should always be true 1597 if module_name[0] == ".": 1598 # strip "." 1599 module_name = module_name[1:] 1600 # If it was a dataparallel it will also have a module at the start 1601 # so strip that for loading 1602 if module_name[:6] == "module": 1603 module_name = module_name[7:] 1604 return module_name 1605 1606 1607def load_net_from_dict(net, state_dict): 1608 """load the network 1609 1610 This is called within load_net 1611 1612 Parameters 1613 ---------- 1614 net : nn.Module 1615 The network to save. 1616 state_dict : dict 1617 The state dictionary to load. 1618 1619 Returns 1620 ------- 1621 nn.Module 1622 The loaded network. 1623 1624 """ 1625 if GPA.pc.get_verbose(): 1626 print("loading net from dict") 1627 pai_modules = get_pai_modules(net, 0) 1628 if pai_modules == []: 1629 print( 1630 "PAI load_net and load_system uses a state_dict so it must be\n" 1631 "called with a net after perforate_model has been called" 1632 ) 1633 print( 1634 "This is being flagged because you are attempting to load a model\n" 1635 "that does not have any pai_modules in it. Confirm that you are calling\n" 1636 "perforate_model on the correct model, and the same model is the one\n" 1637 "being passed into add_validation_score" 1638 ) 1639 import pdb # This needs to be here for cython for some reason. 1640 pdb.set_trace() 1641 sys.exit(-1) 1642 if GPA.pc.get_verbose(): 1643 print( 1644 "setting up arrays and simulating cycles for %d pai modules" 1645 % len(pai_modules) 1646 ) 1647 not_save = GPA.pc.get_module_names_to_not_save() 1648 for module in pai_modules: 1649 if any(module.name.startswith(ns) for ns in not_save): 1650 print("skipping loading %s based on module_names_to_not_save" % module.name) 1651 continue 1652 # Set up name to be what will be saved in the state dict 1653 module_name = get_module_base_name(module) 1654 module.clear_dendrites() 1655 for tracker in module.dendrite_module.dendrite_values: 1656 try: 1657 tracker.setup_arrays( 1658 state_dict[ 1659 module_name + ".dendrite_module.dendrite_values.0.dendrite_storage_shape" 1660 ].tolist() 1661 ) 1662 except Exception as e: 1663 print(e) 1664 first_key = next(iter(state_dict.keys())) 1665 error_msg = f""" 1666This value is missing from the state dict. 1667When missing this value it typically means you converted a module but didn't 1668actually use it in your forward and backward pass. 1669 1670Module was: {module.name} 1671 1672There are many reasons this can happen: 1673 16741 - Check your model definition and forward function and ensure this module is 1675 being used properly. 1676 1677 With GPA.pc.set_verbose(True) you can confirm this is the case if you do not 1678 see a "setting d shape for" this module at the first training batch. 1679 1680 If this is the case, and it is correct to not be passing data through it, 1681 set it to be a tracked module with: 1682 GPA.pc.append_module_ids_to_track(["{module.name}"]) to leave it out 1683 16842 - This can happen if you adjusted your model definition after calling 1685 perforate_model (for example with torch.compile). If the module name printed 1686 above does not contain all modules leading to the main definition, this is 1687 likely the case for your problem. Fix by calling perforate_model after all 1688 other model initialization steps. 1689 16903 - This can happen if the model where you called perforate_model and the model 1691 within add_validation_score are not the same. Check if the module above and 1692 .{first_key} have the same prefix. If one starts with .model or .base etc 1693 and the other does not, this is the problem. 1694 16954 - If you are using this module but then not actually including the correct 1696 output tensor in the forward. For example if you are using an LSTM and 1697 forwarding hidden instead of output but your processors are set up to work 1698 with output. 1699 17005 - If you are not properly calling backward at all. If this is the first module 1701 in your network it is more likely this is the problem. One check in these 1702 cases is to make sure you do not call an initial validation score before the 1703 first backward call. If you do this, while testing_dendrite_capacity is True, 1704 this error will be triggered. 1705 17066 - You have converted a module that is in a frozen part of the network and thus 1707 no gradients are flowing. 1708 17097 - You are running multiple experiments at once with the same save_name. When 1710 running concurrent trials be sure to add save_name=<unique_name> to 1711 perforate_model. 1712""" 1713 print(error_msg) 1714 import pdb # This needs to be here for cython for some reason. 1715 pdb.set_trace() 1716 1717 # Perform as many cycles as the state dict has 1718 num_cycles = int(state_dict[module_name + ".dendrite_module.num_cycles"].item()) 1719 if num_cycles > 0: 1720 simulate_cycles(module, num_cycles, doing_pai=True) 1721 # Handle tracker_string loading with flexible key matching 1722 tracker_key = None 1723 if "tracker_string" in state_dict: 1724 tracker_key = "tracker_string" 1725 else: 1726 # Search for keys containing "tracker_string" 1727 tracker_keys = [key for key in state_dict.keys() if "tracker_string" in key] 1728 if len(tracker_keys) == 1: 1729 tracker_key = tracker_keys[0] 1730 elif len(tracker_keys) > 1: 1731 print(f"Error: Multiple tracker_string keys found: {tracker_keys}") 1732 import pdb # This needs to be here for cython for some reason. 1733 pdb.set_trace() 1734 else: 1735 print("Error: No tracker_string found in state_dict") 1736 import pdb # This needs to be here for cython for some reason. 1737 pdb.set_trace() 1738 1739 if hasattr(net, "tracker_string"): 1740 net.tracker_string = state_dict[tracker_key] 1741 else: 1742 net.register_buffer("tracker_string", state_dict[tracker_key]) 1743 try: 1744 load_result = net.load_state_dict(state_dict, strict=False) 1745 not_save_state_names = [ns.lstrip('.') for ns in not_save] 1746 1747 def is_ignored_key(key): 1748 """Check whether a state-dict key should be ignored. 1749 1750 Parameters 1751 ---------- 1752 key : str 1753 State-dict key to test. 1754 1755 Returns 1756 ------- 1757 bool 1758 ``True`` when key belongs to a not-saved namespace. 1759 """ 1760 return any(key.startswith(ns) for ns in not_save_state_names) 1761 1762 missing_keys = [key for key in load_result.missing_keys if not is_ignored_key(key)] 1763 unexpected_keys = [key for key in load_result.unexpected_keys if not is_ignored_key(key)] 1764 1765 if GPA.pc.get_strict_loading() and (missing_keys or unexpected_keys): 1766 raise RuntimeError( 1767 "Error(s) in loading state_dict for %s:\n\tMissing key(s) in state_dict: %s. \n\tUnexpected key(s) in state_dict: %s." 1768 % (type(net).__name__, missing_keys, unexpected_keys) 1769 ) 1770 except Exception as e: 1771 """ 1772 When modules have high depth to them (i.e. modules within modules not number of layers) 1773 PyTorch can have trouble loading state dicts even when they are correct. 1774 This is a workaround to manually load the state dict if this happens. 1775 """ 1776 filtered_net_keys = { 1777 key 1778 for key in net.state_dict().keys() 1779 if not any(key.startswith(ns.lstrip('.')) for ns in not_save) 1780 } 1781 if filtered_net_keys == set(state_dict.keys()): 1782 print("Attempting manual loading of state_dict") 1783 manual_load_state_dict(net, state_dict) 1784 else: 1785 print(f"Error loading state_dict: {e}") 1786 print("If the error is due to missing keys (e.g., from code changes), you can try:") 1787 print(" GPA.pc.set_strict_loading(False)") 1788 print(" Do not change this unless you are certain the missing keys are not important to load and are expected due to code changes or arch changes.") 1789 print("\ntype 'c' to print full state dicts\n") 1790 import pdb # This needs to be here for cython for some reason. 1791 pdb.set_trace() 1792 print("net state dict is:") 1793 print(net.state_dict()) 1794 print("loaded state dict is:") 1795 print(state_dict) 1796 print( 1797 "Try to check differences. Likely is caused by a module not " 1798 "being converted that should be or vice versa" 1799 ) 1800 pdb.set_trace() 1801 net.to(GPA.pc.get_device()) 1802 return net 1803 1804 1805def pai_save_system(net, folder, name): 1806 """Save the entire system with scaffolding removed 1807 1808 This is used for the final network for inference after training 1809 1810 Parameters 1811 ---------- 1812 net : nn.Module 1813 The network to save. 1814 folder : str 1815 The folder to save the network in. 1816 name : str 1817 The name to save the network under. 1818 1819 Returns 1820 ------- 1821 None 1822 1823 """ 1824 net.member_vars = {} 1825 for member_var in GPA.pai_tracker.member_vars: 1826 if member_var == "scheduler_instance" or member_var == "optimizer_instance": 1827 continue 1828 net.member_vars[member_var] = GPA.pai_tracker.member_vars[member_var] 1829 pai_save_net(net, folder, name) 1830 1831 1832def deep_copy_pai(net): 1833 """Deep copy a PAI network 1834 1835 1836 Parameters 1837 ---------- 1838 net : nn.Module 1839 The network to copy. 1840 1841 Returns 1842 ------- 1843 nn.Module 1844 The copied network. 1845 1846 Notes 1847 ---- 1848 This is required because processors must be cleared before calling copy 1849 1850 """ 1851 # Dont check this stuff if its before the perforate_model has been called and you're just copying a regular model 1852 if(GPA.pai_tracker != []): 1853 # Clear gradients before saving the model 1854 if ((GPA.pai_tracker.member_vars["optimizer_instance"]) is not None) and ( 1855 GPA.pai_tracker.member_vars["optimizer_instance"] != [] 1856 ): 1857 GPA.pai_tracker.member_vars["optimizer_instance"].zero_grad() 1858 GPA.pai_tracker.clear_all_processors() 1859 return copy.deepcopy(net) 1860 1861 1862def prepare_final_model(net): 1863 """Prepare model for final save by removing scaffolding. 1864 1865 This performs all cleanup steps to convert a PAI model with scaffolding 1866 into a clean final model ready for inference or distribution. 1867 1868 Parameters 1869 ---------- 1870 net : nn.Module 1871 The network to prepare. 1872 1873 Returns 1874 ------- 1875 nn.Module 1876 The cleaned model with scaffolding removed. 1877 """ 1878 # Deep copy and clean the model (removes scaffolding) 1879 net = deep_copy_pai(net) 1880 net = BPA.blockwise_network(net) 1881 net = deep_copy_pai(net) 1882 net = CL.refresh_net(net) 1883 1884 # Remove tracker_string (not needed for final model) 1885 if hasattr(net, "tracker_string"): 1886 del net.tracker_string 1887 1888 # Make parameters contiguous 1889 for param in net.parameters(): 1890 param.data = param.data.contiguous() 1891 1892 return net 1893 1894 1895def pai_save_net(net, folder, name): 1896 """Save the entire system with scaffolding removed 1897 1898 This is called within pai_save_system after the tracker has been 1899 turned into a single tensor to be saved as a part of the network 1900 1901 1902 Parameters 1903 ---------- 1904 net : nn.Module 1905 The network to save. 1906 folder : str 1907 The folder to save the network in. 1908 name : str 1909 The name to save the network under. 1910 1911 Returns 1912 ------- 1913 None 1914 1915 Notes 1916 ---- 1917 For open source implementation this is not as important since 1918 minimal values are already being used. 1919 1920 """ 1921 1922 if GPA.pc.get_perforated_backpropagation(): 1923 UPB.pb_save_net(net, folder, name) 1924 else: 1925 return 1926 1927 1928def simulate_cycles(module, num_cycles, doing_pai): 1929 """Simulate dendrite addition cycles 1930 1931 Simulate the back and forth processes of adding dendrites to build a 1932 pretrained dendrite model before loading weights. Required for loading 1933 dendrite save files from non dendrite initial models. 1934 1935 If calling this by hand to start off a network with dendrites added, 1936 make sure to pass a single datapoint forward and backward through the network 1937 so that array sizes are initialized. 1938 1939 Parameters 1940 ---------- 1941 module : PA.PAINeuronModule 1942 The module to simulate cycles on. 1943 num_cycles : int 1944 The number of cycles to simulate. 1945 doing_pai : bool 1946 Whether to actually do the simulation. 1947 1948 Returns 1949 ------- 1950 None 1951 1952 """ 1953 1954 check_skipped = GPA.pc.get_checked_skipped_modules() 1955 if doing_pai is False: 1956 return 1957 GPA.pc.set_checked_skipped_modules(True) 1958 mode = "n" 1959 for i in range(num_cycles): 1960 if mode == "n": 1961 module.set_mode("p") 1962 module.create_new_dendrite_module() 1963 mode = "p" 1964 else: 1965 module.set_mode("n") 1966 mode = "n" 1967 GPA.pc.set_checked_skipped_modules(check_skipped) 1968 1969 1970def count_params(net): 1971 """Count the number of parameters in the network 1972 1973 If doing perforated backpropagation this calls the PB function 1974 which does not count scaffolding parameters since the final model 1975 will not have them. 1976 1977 Parameters 1978 ---------- 1979 net : nn.Module 1980 The network to count parameters in. 1981 1982 Returns 1983 ------- 1984 int 1985 The number of parameters in the network. 1986 1987 """ 1988 if GPA.pc.get_perforated_backpropagation(): 1989 return UPB.pb_count_params(net) 1990 parameters = net.named_parameters() 1991 unique_params = { 1992 p.data_ptr(): p for name, p in parameters if "parent_module" not in name 1993 }.values() 1994 return sum(p.numel() for p in unique_params) 1995 1996 1997def change_learning_modes(net, folder, name, doing_pai): 1998 """Change between neuron and dendrite learning modes 1999 2000 High level steps for entire system to switch back and forth between 2001 neuron learning and dendrite learning 2002 2003 Parameters 2004 ---------- 2005 net : nn.Module 2006 The network to change modes on. 2007 folder : str 2008 The folder to save/load the network in/from. 2009 name : str 2010 The name to save/load the network under. 2011 doing_pai : bool 2012 Whether to add dendrites when changing modes. 2013 2014 Returns 2015 ------- 2016 int 2017 The number of parameters in the network. 2018 2019 Notes 2020 ----- 2021 If doing_pai is False this just allows training to continue longer rather than early stopping 2022 2023 """ 2024 # If not adding dendrites this just allows training to continue longer with flags 2025 # every time early stopping should be occurring 2026 if doing_pai is False: 2027 GPA.pai_tracker.member_vars["switch_epochs"].append( 2028 GPA.pai_tracker.member_vars["num_epochs_run"] 2029 ) 2030 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2031 "switch_epochs" 2032 ][-1] 2033 GPA.pai_tracker.reset_vals_for_score_reset() 2034 return net 2035 if GPA.pai_tracker.member_vars["mode"] == "n": 2036 current_epoch = GPA.pai_tracker.member_vars["num_epochs_run"] 2037 overwritten_epochs = GPA.pai_tracker.member_vars["overwritten_epochs"] 2038 overwritten_extra = GPA.pai_tracker.member_vars["extra_scores"] 2039 if GPA.pc.get_drawing_pai(): 2040 overwritten_val = GPA.pai_tracker.member_vars["accuracies"] 2041 else: 2042 overwritten_val = GPA.pai_tracker.member_vars["neuron_accuracies"] 2043 """ 2044 If true don't load the best system 2045 because it will delete dendrites if the previous best was better than 2046 the current best 2047 """ 2048 if not GPA.pc.get_silent(): 2049 print("Importing best Model for switch to PA...") 2050 net = load_system(net, folder, name, switch_call=True) 2051 GPA.pai_tracker.set_dendrite_training() 2052 GPA.pai_tracker.member_vars["overwritten_epochs"] = overwritten_epochs 2053 GPA.pai_tracker.member_vars["overwritten_epochs"] += ( 2054 current_epoch - GPA.pai_tracker.member_vars["num_epochs_run"] 2055 ) 2056 GPA.pai_tracker.member_vars["total_epochs_run"] = ( 2057 GPA.pai_tracker.member_vars["num_epochs_run"] 2058 + GPA.pai_tracker.member_vars["overwritten_epochs"] 2059 ) 2060 2061 if GPA.pc.get_save_old_graph_scores(): 2062 GPA.pai_tracker.member_vars["overwritten_extras"].append(overwritten_extra) 2063 GPA.pai_tracker.member_vars["overwritten_vals"].append(overwritten_val) 2064 else: 2065 GPA.pai_tracker.member_vars["overwritten_extras"] = [overwritten_extra] 2066 GPA.pai_tracker.member_vars["overwritten_vals"] = [overwritten_val] 2067 if GPA.pc.get_drawing_pai(): 2068 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2069 GPA.pai_tracker.member_vars["num_epochs_run"] 2070 ) 2071 else: 2072 if len(GPA.pai_tracker.member_vars["switch_epochs"]) == 0: 2073 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2074 GPA.pai_tracker.member_vars["num_epochs_run"] 2075 ) 2076 else: 2077 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2078 GPA.pai_tracker.member_vars["n_switch_epochs"][-1] 2079 + ( 2080 (GPA.pai_tracker.member_vars["num_epochs_run"]) 2081 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2082 ) 2083 ) 2084 2085 GPA.pai_tracker.member_vars["switch_epochs"].append( 2086 GPA.pai_tracker.member_vars["num_epochs_run"] 2087 ) 2088 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2089 "switch_epochs" 2090 ][-1] 2091 2092 # Because open source version is only doing neuron training for 2093 # gradient descent dendrites, switch back to n mode right away 2094 if ( 2095 not GPA.pc.get_perforated_backpropagation() 2096 ) or GPA.pc.get_no_extra_n_modes(): 2097 net = change_learning_modes(net, folder, name, doing_pai) 2098 else: 2099 if not GPA.pc.get_silent(): 2100 print("Switching back to N...") 2101 set_best = GPA.pai_tracker.member_vars["current_n_set_global_best"] 2102 GPA.pai_tracker.set_neuron_training() 2103 if len(GPA.pai_tracker.member_vars["p_switch_epochs"]) == 0: 2104 GPA.pai_tracker.member_vars["p_switch_epochs"].append( 2105 ( 2106 (GPA.pai_tracker.member_vars["num_epochs_run"] - 1) 2107 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2108 ) 2109 ) 2110 else: 2111 GPA.pai_tracker.member_vars["p_switch_epochs"].append( 2112 GPA.pai_tracker.member_vars["p_switch_epochs"][-1] 2113 + ( 2114 (GPA.pai_tracker.member_vars["num_epochs_run"]) 2115 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2116 ) 2117 ) 2118 GPA.pai_tracker.member_vars["switch_epochs"].append( 2119 GPA.pai_tracker.member_vars["num_epochs_run"] 2120 ) 2121 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2122 "switch_epochs" 2123 ][-1] 2124 # Will be false for open source implementation 2125 if GPA.pc.get_retain_all_dendrites() or ( 2126 GPA.pc.get_learn_dendrites_live() and set_best 2127 ): 2128 if not GPA.pc.get_silent(): 2129 print( 2130 "Saving model before starting normal training to " 2131 "retain PBNodes regardless of next N Phase results" 2132 ) 2133 save_system(net, folder, name) 2134 # if its just doing P for learn PAI live then switch back immediately 2135 if GPA.pc.get_perforated_backpropagation() and GPA.pc.get_no_extra_n_modes(): 2136 net = change_learning_modes(net, folder, name, doing_pai) 2137 2138 GPA.pai_tracker.member_vars["param_counts"].append(count_params(net)) 2139 2140 return net 2141 2142 2143def find_param_name_by_id(model, param_id): 2144 """ 2145 This is only used for debugging. 2146 Return the fully-qualified parameter name (e.g. "layer1.conv.weight") 2147 for the parameter whose id matches param_id. Returns None if not found. 2148 2149 This uses model.named_parameters(), which already recurses through submodules. 2150 2151 Parameters 2152 ---------- 2153 model : Model to look for param id. 2154 param_id : pointer to a parameter. 2155 2156 Returns 2157 ------- 2158 String representing the module within the model or None if not found. 2159 """ 2160 for name, p in model.named_parameters(recurse=True): 2161 if id(p) == param_id: 2162 return "." + name 2163 return None 2164 2165 2166def add_method_delegation_to_module(wrapper_module, method_name): 2167 """Add delegating methods to a wrapper module that has a main_module attribute. 2168 2169 This adds the specified methods to the wrapper module instance so they 2170 properly delegate to the wrapped main_module. Works for any wrapper module 2171 (TrackedNeuronModule, PAINeuronModule, etc.) that has a main_module attribute. 2172 2173 Args: 2174 wrapper_module: A wrapper module instance with a main_module attribute 2175 method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') 2176 2177 Parameters 2178 ---------- 2179 wrapper_module : PyTorch Module that contains a sub module. 2180 method_name : Name of method to wrap. 2181 2182 Returns 2183 ------- 2184 None 2185 """ 2186 import types 2187 2188 if hasattr(wrapper_module.main_module, method_name): 2189 # Create a delegating method that forwards to main_module 2190 def make_delegated_method(name): 2191 """Create a bound delegation function for a given attribute name. 2192 2193 Parameters 2194 ---------- 2195 name : str 2196 Attribute name to forward to ``main_module``. 2197 2198 Returns 2199 ------- 2200 callable 2201 Function that delegates access or invocation. 2202 """ 2203 def delegated_method(self, *args, **kwargs): 2204 """Delegate attribute access or method call to ``main_module``. 2205 2206 Parameters 2207 ---------- 2208 *args : tuple 2209 Positional arguments forwarded to delegated callables. 2210 **kwargs : dict 2211 Keyword arguments forwarded to delegated callables. 2212 2213 Returns 2214 ------- 2215 Any 2216 Delegated attribute value or method result. 2217 """ 2218 main_module_attr = getattr(self.main_module, name, None) 2219 if main_module_attr is None: 2220 raise AttributeError( 2221 f"'{type(self.main_module).__name__}' object has no attribute '{name}'" 2222 ) 2223 if callable(main_module_attr): 2224 return main_module_attr(*args, **kwargs) 2225 return main_module_attr 2226 2227 return delegated_method 2228 2229 # Bind it to this specific instance 2230 setattr( 2231 wrapper_module, 2232 method_name, 2233 types.MethodType(make_delegated_method(method_name), wrapper_module), 2234 ) 2235 2236 2237def apply_method_delegation_to_model(model, method_name, main_module_type): 2238 """Recursively apply method delegation to all wrapper modules with main_module in a model. 2239 2240 This traverses the entire model and adds method delegation for any module that has 2241 a main_module attribute and optionally matches specified types. 2242 2243 Args: 2244 model: The PyTorch model to traverse 2245 method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') 2246 main_module_type: main_module type name to filter by. 2247 Example: 'Qwen2DecoderLayer' 2248 2249 Example: 2250 # Apply gradient checkpointing delegation to all decoder layers 2251 apply_method_delegation_to_model( 2252 model, 2253 '_gradient_checkpointing_func', 2254 main_module_type='Qwen2DecoderLayer' 2255 ) 2256 2257 Parameters 2258 ---------- 2259 model : PyTorch model. 2260 method_name : method to delegate. 2261 main_module_type : type of module that has this method. 2262 2263 Returns 2264 ------- 2265 None 2266 This function does not return a value. 2267 """ 2268 count = 0 2269 for name, module in model.named_modules(): 2270 # Check if module has main_module attribute (it's a wrapper) 2271 if hasattr(module, "main_module"): 2272 # Check if we should apply based on main_module type 2273 should_apply = True 2274 if main_module_type is not None: 2275 main_module_type_name = type(module.main_module).__name__ 2276 should_apply = main_module_type_name == main_module_type 2277 2278 if should_apply: 2279 add_method_delegation_to_module(module, method_name) 2280 count += 1 2281 2282 print(f"[PAI] Applied method delegation to {count} wrapper module instances") 2283 2284 2285def make_json_serializable(obj): 2286 """Recursively convert non-JSON-serializable objects to strings. 2287 2288 Parameters 2289 ---------- 2290 obj : any 2291 The object to convert 2292 2293 Returns 2294 ------- 2295 Any 2296 JSON-serializable version of the object 2297 """ 2298 if isinstance(obj, (str, int, float, bool, type(None))): 2299 return obj 2300 elif isinstance(obj, dict): 2301 return {k: make_json_serializable(v) for k, v in obj.items()} 2302 elif isinstance(obj, (list, tuple)): 2303 return [make_json_serializable(item) for item in obj] 2304 else: 2305 # Convert non-serializable types to string 2306 return str(obj) 2307 2308 2309def extract_gpa_config(): 2310 """Extract all configuration from GPA.pc by calling all get_* methods. 2311 2312 Returns 2313 ------- 2314 dict[str, Any] 2315 Dictionary with all GPA.pc configuration values and type metadata 2316 2317 Examples 2318 -------- 2319 >>> config = extract_gpa_config() 2320 >>> # Returns: {'max_dendrites': 10, 'device': 'cuda', '_types': {...}} 2321 2322 Parameters 2323 ---------- 2324 None 2325 2326 """ 2327 config = {} 2328 config_types = {} 2329 2330 # Get all attributes from GPA.pc 2331 for attr_name in dir(GPA.pc): 2332 # Check if it starts with 'get_' 2333 if attr_name.startswith("get_"): 2334 try: 2335 # Get the method 2336 method = getattr(GPA.pc, attr_name) 2337 2338 # Check if it's callable 2339 if callable(method): 2340 # Call it and store result with key as name without 'get_' 2341 key = attr_name[4:] # Remove 'get_' prefix 2342 value = method() 2343 2344 # Check if this is an array (has corresponding append_ method) 2345 append_method_name = f"append_{key}" 2346 is_array = hasattr(GPA.pc, append_method_name) 2347 2348 if is_array and isinstance(value, (list, tuple)): 2349 # Store array element type 2350 if len(value) > 0: 2351 element_type = type(value[0]).__name__ 2352 else: 2353 element_type = None # empty array, no conversion needed 2354 config_types[key] = { 2355 "is_array": True, 2356 "element_type": element_type, 2357 } 2358 else: 2359 # Store value type 2360 config_types[key] = { 2361 "is_array": False, 2362 "type": type(value).__name__, 2363 } 2364 2365 # Make sure value is JSON serializable 2366 config[key] = make_json_serializable(value) 2367 except Exception as e: 2368 # Skip if method fails 2369 if GPA.pc.get_verbose(): 2370 print(f"Skipping {attr_name}: {e}") 2371 continue 2372 2373 # Add types metadata to config 2374 config["_types"] = config_types 2375 2376 return config 2377 2378 2379def convert_to_type(value, type_name): 2380 """Convert a value to the specified type. 2381 2382 Parameters 2383 ---------- 2384 value : any 2385 The value to convert 2386 type_name : str 2387 The target type name 2388 2389 Returns 2390 ------- 2391 Any 2392 The converted value 2393 """ 2394 if type_name == "NoneType" or value is None: 2395 return None 2396 elif type_name == "bool": 2397 if isinstance(value, str): 2398 return value.lower() in ("true", "1", "yes") 2399 return bool(value) 2400 elif type_name == "int": 2401 return int(value) 2402 elif type_name == "float": 2403 return float(value) 2404 elif type_name == "str": 2405 return str(value) 2406 elif type_name == "list": 2407 if not isinstance(value, list): 2408 return [value] 2409 return value 2410 elif type_name == "dict": 2411 if not isinstance(value, dict): 2412 return {} 2413 return value 2414 elif type_name == "type": 2415 # Handle type objects - convert string representation back to type 2416 if isinstance(value, str): 2417 # Try to evaluate the type string (e.g., "<class 'torch.nn.Linear'>") 2418 # Extract the class path from the string 2419 if value.startswith("<class '") and value.endswith("'>"): 2420 class_path = value[ 2421 8:-2 2422 ] # Extract 'torch.nn.Linear' from "<class 'torch.nn.Linear'>" 2423 parts = class_path.split(".") 2424 # Try to import and get the type 2425 try: 2426 module_name = ".".join(parts[:-1]) 2427 class_name = parts[-1] 2428 module = __import__(module_name, fromlist=[class_name]) 2429 return getattr(module, class_name) 2430 except Exception as e: 2431 print( 2432 f"Warning: Could not convert type string '{value}' to actual type: {e}" 2433 ) 2434 return value 2435 return value 2436 return value 2437 elif type_name == "dtype": 2438 # Handle torch dtype objects 2439 if isinstance(value, str): 2440 # Convert string like "torch.float32" to actual dtype 2441 import torch 2442 2443 try: 2444 # Try to get the dtype from torch module 2445 if value.startswith("torch."): 2446 dtype_name = value.split(".")[ 2447 1 2448 ] # Get 'float32' from 'torch.float32' 2449 return getattr(torch, dtype_name) 2450 else: 2451 return getattr(torch, value) 2452 except Exception as e: 2453 print( 2454 f"Warning: Could not convert dtype string '{value}' to actual dtype: {e}" 2455 ) 2456 return value 2457 return value 2458 elif type_name == "device": 2459 # Handle torch device objects 2460 if isinstance(value, str): 2461 # Convert string like "cuda" or "cpu" to torch.device 2462 import torch 2463 2464 try: 2465 return torch.device(value) 2466 except Exception as e: 2467 print( 2468 f"Warning: Could not convert device string '{value}' to actual device: {e}" 2469 ) 2470 return value 2471 return value 2472 elif type_name == "builtin_function_or_method": 2473 # Handle torch functions like torch.sigmoid, torch.relu, etc. 2474 if isinstance(value, str): 2475 # Parse string like "<built-in method sigmoid of type object at 0x...>" 2476 # to extract the function name 2477 import torch 2478 2479 try: 2480 if "<built-in method " in value and " of type object" in value: 2481 # Extract function name between '<built-in method ' and ' of type object' 2482 start = value.find("<built-in method ") + len("<built-in method ") 2483 end = value.find(" of type object") 2484 func_name = value[start:end] 2485 # Try to get the function from torch module 2486 if hasattr(torch, func_name): 2487 return getattr(torch, func_name) 2488 else: 2489 print(f"Warning: torch.{func_name} not found") 2490 return value 2491 else: 2492 return value 2493 except Exception as e: 2494 print( 2495 f"Warning: Could not convert builtin function string '{value}': {e}" 2496 ) 2497 return value 2498 return value 2499 else: 2500 # Unknown type - error and debug 2501 print(f"ERROR: Unknown type '{type_name}' for value: {value}") 2502 print(f"Type of value is: {type(value).__name__}") 2503 pdb.set_trace() 2504 return value 2505 2506 2507def convert_to_type_array(value, element_type): 2508 """Convert an array's elements to the specified type. 2509 2510 Parameters 2511 ---------- 2512 value : list or tuple 2513 The array to convert 2514 element_type : str or None 2515 The target type name for elements, None if array was empty 2516 2517 Returns 2518 ------- 2519 list 2520 The array with converted elements 2521 """ 2522 if not isinstance(value, (list, tuple)): 2523 return value 2524 # If element_type is None (empty array), no conversion needed 2525 if element_type is None: 2526 return list(value) if isinstance(value, tuple) else value 2527 return [convert_to_type(item, element_type) for item in value] 2528 2529 2530def set_gpa_config(config): 2531 """Set GPA.pc configuration by calling all set_* methods. 2532 2533 This is the reverse of extract_gpa_config(). It takes a configuration 2534 dictionary and calls the corresponding set_* methods on GPA.pc. 2535 Uses type metadata to ensure values are converted to the correct type. 2536 2537 Parameters 2538 ---------- 2539 config : dict 2540 Dictionary with configuration values (keys without 'set_' prefix) 2541 and optional '_types' metadata 2542 2543 Examples 2544 -------- 2545 >>> config = {'verbose': True, 'device': 'cuda'} 2546 >>> set_gpa_config(config) 2547 # Calls GPA.pc.set_verbose(True), GPA.pc.set_device('cuda'), etc. 2548 2549 Returns 2550 ------- 2551 Count of parameters that were set 2552 """ 2553 set_count = 0 2554 skip_count = 0 2555 2556 # Extract type information 2557 config_types = config.get("_types", {}) 2558 2559 for key, value in config.items(): 2560 # Skip the types metadata 2561 if key == "_types": 2562 continue 2563 2564 # Construct the set method name 2565 set_method_name = f"set_{key}" 2566 2567 # Check if the set method exists 2568 if hasattr(GPA.pc, set_method_name): 2569 try: 2570 method = getattr(GPA.pc, set_method_name) 2571 if callable(method): 2572 # Convert value to correct type if we have type info 2573 if key in config_types: 2574 type_info = config_types[key] 2575 if type_info.get("is_array", False): 2576 # Convert array elements to correct type 2577 element_type = type_info.get("element_type", "str") 2578 value = convert_to_type_array(value, element_type) 2579 else: 2580 # Convert single value to correct type 2581 value_type = type_info.get("type", "str") 2582 value = convert_to_type(value, value_type) 2583 2584 method(value) 2585 set_count += 1 2586 if GPA.pc.get_verbose(): 2587 print(f"Set {key} = {value}") 2588 except Exception as e: 2589 skip_count += 1 2590 if GPA.pc.get_verbose(): 2591 print(f"Failed to set {key}: {e}") 2592 else: 2593 skip_count += 1 2594 if GPA.pc.get_verbose(): 2595 print(f"No setter found for {key} (looking for {set_method_name})") 2596 2597 if GPA.pc.get_verbose(): 2598 print(f"Applied {set_count} PAI configuration settings ({skip_count} skipped)") 2599 2600 return set_count 2601 2602 2603try: 2604 from huggingface_hub import PyTorchModelHubMixin, hf_hub_download, HfApi 2605 2606 def upload_to_huggingface( 2607 model, 2608 repo_id, 2609 license="apache-2.0", 2610 pipeline_tag=None, 2611 repo_url=None, 2612 tags=None, 2613 include_pai_config=True, 2614 **kwargs, 2615 ): 2616 """Upload a model to HuggingFace Hub. 2617 2618 Uploads model weights and PAI configuration to HuggingFace Hub. 2619 The configuration is saved in config.json and can be restored when loading. 2620 2621 Parameters 2622 ---------- 2623 model : nn.Module 2624 The model to upload 2625 repo_id : str 2626 Repository ID (format: "username/model-name") 2627 license : str, optional 2628 License for the model card, by default "apache-2.0" 2629 pipeline_tag : str, optional 2630 Pipeline tag for the model (e.g., "text-classification", "image-classification") 2631 repo_url : str, optional 2632 URL to the model's repository/documentation 2633 tags : list, optional 2634 List of tags for the model card 2635 include_pai_config : bool, optional 2636 Whether to include all GPA.pc configuration in the model config, by default True 2637 **kwargs 2638 Additional arguments passed to HfApi (token, private, etc.) 2639 2640 Returns 2641 ------- 2642 str 2643 URL of the uploaded model 2644 2645 Examples 2646 -------- 2647 >>> url = upload_to_huggingface( 2648 ... model, 2649 ... "username/my-model", 2650 ... license="mit", 2651 ... pipeline_tag="image-classification", 2652 ... tags=["pytorch", "vision"] 2653 ... ) 2654 """ 2655 try: 2656 from huggingface_hub import HfApi 2657 except ImportError: 2658 raise ImportError( 2659 "huggingface_hub is required. Install it with: pip install huggingface_hub" 2660 ) 2661 2662 import tempfile 2663 import os 2664 2665 # Prepare model same way as save_pai_net does 2666 model = prepare_final_model(model) 2667 2668 # Calculate parameter count 2669 param_count = count_params(model) 2670 2671 # Format parameter count for tags (e.g., "11m" for 11 million) 2672 if param_count >= 1e9: 2673 param_tag = f"{param_count/1e9:.0f}b" 2674 elif param_count >= 1e6: 2675 param_tag = f"{param_count/1e6:.0f}m" 2676 elif param_count >= 1e3: 2677 param_tag = f"{param_count/1e3:.0f}k" 2678 else: 2679 param_tag = str(param_count) 2680 2681 # Create a temporary directory for files 2682 with tempfile.TemporaryDirectory() as tmpdir: 2683 # Save model weights 2684 model_path = os.path.join(tmpdir, "model.safetensors") 2685 save_file(model.state_dict(), model_path) 2686 2687 # Create config with PAI configuration 2688 config = {} 2689 if include_pai_config: 2690 pai_config = extract_gpa_config() 2691 config["pai_config"] = pai_config 2692 if GPA.pc.get_verbose(): 2693 print(f"Extracted {len(pai_config)} PAI configuration parameters") 2694 2695 # Add parameter count at top level 2696 config["num_parameters"] = param_count 2697 2698 # Add metadata 2699 if license: 2700 config["license"] = license 2701 if pipeline_tag: 2702 config["pipeline_tag"] = pipeline_tag 2703 if repo_url: 2704 config["repo_url"] = repo_url 2705 2706 # Add tags with parameter count 2707 if tags is None: 2708 tags = [] 2709 elif not isinstance(tags, list): 2710 tags = [tags] 2711 else: 2712 tags = tags.copy() # Don't modify the original list 2713 2714 # Add perforated-ai tag if not present 2715 if "perforated-ai" not in tags: 2716 tags.insert(0, "perforated-ai") 2717 2718 # Add parameter count tag if not present 2719 if param_tag not in tags: 2720 tags.append(param_tag) 2721 2722 config["tags"] = tags 2723 2724 # Save config.json 2725 config_path = os.path.join(tmpdir, "config.json") 2726 with open(config_path, "w") as f: 2727 json.dump(config, f, indent=2) 2728 2729 # Upload to HuggingFace 2730 api = HfApi() 2731 2732 # Extract token from kwargs if present 2733 token = kwargs.pop("token", None) 2734 private = kwargs.pop("private", None) 2735 2736 # Create repo if it doesn't exist 2737 try: 2738 api.create_repo( 2739 repo_id=repo_id, token=token, private=private, exist_ok=True 2740 ) 2741 except Exception as e: 2742 print(f"Repo may already exist: {e}") 2743 2744 # Upload folder 2745 api.upload_folder( 2746 folder_path=tmpdir, repo_id=repo_id, token=token, **kwargs 2747 ) 2748 2749 print(f"Model uploaded to: https://huggingface.co/{repo_id}") 2750 if include_pai_config: 2751 print(f"PAI configuration saved in config.json") 2752 print(f"To reload, use: model = from_hf_pretrained(model, '{repo_id}')") 2753 2754 return f"https://huggingface.co/{repo_id}" 2755 2756 def from_hf_pretrained(net, repo_id, force_download=False): 2757 """Load a PerforatedAI model from HuggingFace Hub using PyTorchModelHubMixin. 2758 2759 Args: 2760 net: The base model architecture (will be converted to PAI format) 2761 repo_id: HuggingFace Hub repository ID (e.g., "username/model-name") 2762 force_download: If True, always download the latest version, bypassing cache (default: False) 2763 2764 Returns: 2765 net: The loaded model with PAI modules initialized 2766 2767 Parameters 2768 ---------- 2769 net : PyTorch Model. 2770 repo_id : Name of HuggingFace repository. 2771 force_download : Force an update even if local file exists. 2772 2773 Returns 2774 ------- 2775 Loaded Model 2776 """ 2777 2778 # Wrap in a class that inherits from PyTorchModelHubMixin 2779 class PAIHFModel(net.__class__, PyTorchModelHubMixin): 2780 def __init__(self, *args, **kwargs): 2781 """Initialize the temporary HuggingFace-compatible wrapper.""" 2782 super().__init__(*args, **kwargs) 2783 2784 # Create an instance that can use from_pretrained 2785 wrapped_net = PAIHFModel.__new__(PAIHFModel) 2786 wrapped_net.__dict__ = net.__dict__ 2787 wrapped_net.__class__ = PAIHFModel 2788 2789 # Download config.json to restore PAI configuration 2790 try: 2791 config_path = hf_hub_download(repo_id=repo_id, filename="config.json", force_download=force_download) 2792 with open(config_path, "r") as f: 2793 config = json.load(f) 2794 if "pai_config" in config: 2795 # print(f"Restoring PAI configuration from HuggingFace") 2796 set_gpa_config(config["pai_config"]) 2797 else: 2798 print("Warning: No pai_config found in config.json") 2799 except Exception as e: 2800 print(f"Warning: Could not load PAI config from HuggingFace: {e}") 2801 2802 # Download model files from HuggingFace 2803 model_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", force_download=force_download) 2804 state_dict = load_file(model_path) 2805 wrapped_net = NPA.load_pai_model_from_dict(wrapped_net, state_dict) 2806 return wrapped_net 2807 2808except: 2809 2810 def upload_to_huggingface(*args, **kwargs): 2811 """Raise an informative error when HuggingFace dependencies are missing. 2812 2813 Parameters 2814 ---------- 2815 *args : tuple[Any, ...] 2816 Positional arguments accepted for API compatibility. 2817 **kwargs : dict[str, Any] 2818 Keyword arguments accepted for API compatibility. 2819 2820 Returns 2821 ------- 2822 None 2823 Always raises ``ImportError``. 2824 """ 2825 raise ImportError( 2826 "huggingface_hub is required for upload_to_huggingface. " 2827 "Install it with: pip install huggingface_hub" 2828 ) 2829 2830 def from_hf_pretrained(*args, **kwargs): 2831 """Raise an informative error when HuggingFace dependencies are missing. 2832 2833 Parameters 2834 ---------- 2835 *args : tuple[Any, ...] 2836 Positional arguments accepted for API compatibility. 2837 **kwargs : dict[str, Any] 2838 Keyword arguments accepted for API compatibility. 2839 2840 Returns 2841 ------- 2842 None 2843 Always raises ``ImportError``. 2844 """ 2845 raise ImportError( 2846 "huggingface_hub is required for from_hf_pretrained. " 2847 "Install it with: pip install huggingface_hub" 2848 )
def
perforate_model( model, doing_pai=True, save_name='', making_graphs=True, maximizing_score=None, num_classes=10000000000, values_per_train_epoch=-1, values_per_val_epoch=-1, zooming_graph=True, config_file=None):
48def perforate_model( 49 model, 50 doing_pai=True, 51 save_name="", 52 making_graphs=True, 53 maximizing_score=None, 54 num_classes=10000000000, 55 values_per_train_epoch=-1, 56 values_per_val_epoch=-1, 57 zooming_graph=True, 58 config_file=None, 59): 60 """Main function to initialize the network to add dendrites 61 62 This kicks off the entire Perforated AI process to add 63 the scaffolding to the network to be able to add dendrites 64 65 Parameters 66 ---------- 67 model : nn.Module 68 The neural network model to initialize. 69 doing_pai : bool, optional 70 Whether to actually add dendrites, by default True 71 save_name : str, optional 72 The name to save the model under, by default "PAI" 73 making_graphs : bool, optional 74 Whether to create graphs during training, by default True 75 maximizing_score : bool or None, optional 76 Whether to maximize the score during training. 77 - None: use config value (defaults to True) 78 - bool: override config value for this run 79 Setting False is for when the score is a loss to be minimized. 80 num_classes : int, optional 81 The number of output classes, unused in current version 82 values_per_train_epoch : int, optional 83 The number of values to look back for graphing 84 during training, by default -1 (all values). 85 values_per_val_epoch : int, optional 86 The number of values to look back for graphing 87 during validation, by default -1 (all values). 88 zooming_graph : bool, optional 89 Whether to enable zooming on the graphs, by default True 90 config_file : str or None, optional 91 Optional local JSON config file path to load, by default None. 92 93 Returns 94 ------- 95 model : nn.Module 96 The modified model with dendrite scaffolding added if doing_pai is True 97 98 """ 99 100 if save_name == "": 101 if GPA.pc.get_save_name() == "": 102 save_name = "PAI" 103 else: 104 save_name = GPA.pc.get_save_name() 105 106 if "/" in save_name: 107 print( 108 f"Warning: save_name '{save_name}' contains '/'. Relative paths are not implemented yet." 109 ) 110 sys.exit(1) 111 112 sanitized_save_name = "".join( 113 ch for ch in save_name if ch.isalnum() or ch in ("_", "-", ".") 114 ) 115 if sanitized_save_name != save_name: 116 print( 117 f"Warning: save_name '{save_name}' contained spaces or special characters. " 118 f"Using '{sanitized_save_name}' instead." 119 ) 120 save_name = sanitized_save_name 121 122 if save_name == "": 123 print("Warning: save_name became empty after sanitization. Using 'PAI'.") 124 save_name = "PAI" 125 126 if config_file is not None: 127 GPA.pc.set_config_file(config_file) 128 129 GPA.pc.set_save_name(save_name) 130 GPA.pc.sync_config_sources() 131 132 if maximizing_score is None: 133 maximizing_score = GPA.pc.get_maximizing_score() 134 else: 135 # Programmatic perforate_model input must win over config-file values. 136 GPA.pc.set_maximizing_score(maximizing_score) 137 138 if not GPA.pc.get_configuration_confirmed(): 139 CPA.set_perforation_targets(model) 140 141 142 GPA.pai_tracker = TPA.PAINeuronModuleTracker( 143 doing_pai=doing_pai, save_name=save_name 144 ) 145 if _dashboard_emitter is not None: 146 _dashboard_emitter.emit_run_start(GPA.pc, save_name) 147 148 # Get config outputs before tracker is initialized 149 GPA.pc.persist_config_outputs(overwrite_config_file=False) 150 151 model = GPA.pai_tracker.initialize( 152 model, 153 doing_pai=doing_pai, 154 save_name=save_name, 155 making_graphs=making_graphs, 156 maximizing_score=maximizing_score, 157 num_classes=num_classes, 158 values_per_train_epoch=-values_per_train_epoch, 159 values_per_val_epoch=values_per_val_epoch, 160 zooming_graph=zooming_graph, 161 ) 162 163 164 165 return model
Main function to initialize the network to add dendrites
This kicks off the entire Perforated AI process to add the scaffolding to the network to be able to add dendrites
Parameters
- model (nn.Module): The neural network model to initialize.
- doing_pai (bool, optional): Whether to actually add dendrites, by default True
- save_name (str, optional): The name to save the model under, by default "PAI"
- making_graphs (bool, optional): Whether to create graphs during training, by default True
- maximizing_score (bool or None, optional):
Whether to maximize the score during training.
- None: use config value (defaults to True)
- bool: override config value for this run Setting False is for when the score is a loss to be minimized.
- num_classes (int, optional): The number of output classes, unused in current version
- values_per_train_epoch (int, optional): The number of values to look back for graphing during training, by default -1 (all values).
- values_per_val_epoch (int, optional): The number of values to look back for graphing during validation, by default -1 (all values).
- zooming_graph (bool, optional): Whether to enable zooming on the graphs, by default True
- config_file (str or None, optional): Optional local JSON config file path to load, by default None.
Returns
- model (nn.Module): The modified model with dendrite scaffolding added if doing_pai is True
def
get_pai_modules(net, depth, seen_ids=None):
168def get_pai_modules(net, depth, seen_ids=None): 169 """Get a list of all neuron modules 170 171 Parameters 172 ---------- 173 net : nn.Module 174 The module to search. 175 depth : int 176 The current depth in the recursion. 177 178 Returns 179 ------- 180 list 181 A list of all PAI neuron modules found in the network. 182 183 """ 184 if seen_ids is None: 185 seen_ids = set() 186 all_members = net.__dir__() 187 this_list = [] 188 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 189 for submodule_id, layer in net.named_children(): 190 # If there is a self pointer ignore it 191 if net.get_submodule(submodule_id) is net: 192 continue 193 if type(net.get_submodule(submodule_id)) is PA.PAINeuronModule: 194 module = net.get_submodule(submodule_id) 195 if id(module) in seen_ids: 196 continue 197 seen_ids.add(id(module)) 198 this_list = this_list + [module] 199 else: 200 this_list = this_list + get_pai_modules( 201 net.get_submodule(submodule_id), depth + 1, seen_ids 202 ) 203 else: 204 for member in all_members: 205 if isinstance(getattr(type(net), member, None), property): 206 continue 207 # if the getter fails or it is a self pointer ignore it 208 try: 209 if getattr(net, member, None) is net: 210 continue 211 except: 212 continue 213 if type(getattr(net, member, None)) is PA.PAINeuronModule: 214 module = getattr(net, member) 215 if id(module) in seen_ids: 216 continue 217 seen_ids.add(id(module)) 218 this_list = this_list + [module] 219 elif ( 220 issubclass(type(getattr(net, member, None)), nn.Module) 221 or issubclass(type(getattr(net, member, None)), nn.Sequential) 222 or issubclass(type(getattr(net, member, None)), nn.ModuleList) 223 ): 224 this_list = this_list + get_pai_modules( 225 getattr(net, member), depth + 1, seen_ids 226 ) 227 228 return this_list
Get a list of all neuron modules
Parameters
- net (nn.Module): The module to search.
- depth (int): The current depth in the recursion.
Returns
- list: A list of all PAI neuron modules found in the network.
def
get_tracked_modules(net, depth, seen_ids=None):
231def get_tracked_modules(net, depth, seen_ids=None): 232 """Get a list of all tracked modules 233 234 Parameters 235 ---------- 236 net : nn.Module 237 The module to search. 238 depth : int 239 The current depth in the recursion. 240 241 Returns 242 ------- 243 list 244 A list of all tracked modules found in the network. 245 246 """ 247 if seen_ids is None: 248 seen_ids = set() 249 all_members = net.__dir__() 250 this_list = [] 251 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 252 for submodule_id, layer in net.named_children(): 253 if net.get_submodule(submodule_id) is net: 254 continue 255 if type(net.get_submodule(submodule_id)) is PA.TrackedNeuronModule: 256 module = net.get_submodule(submodule_id) 257 if id(module) in seen_ids: 258 continue 259 seen_ids.add(id(module)) 260 this_list = this_list + [module] 261 else: 262 this_list = this_list + get_tracked_modules( 263 net.get_submodule(submodule_id), depth + 1, seen_ids 264 ) 265 else: 266 for member in all_members: 267 if isinstance(getattr(type(net), member, None), property): 268 continue 269 # if the getter fails or it is a self pointer ignore it 270 try: 271 if getattr(net, member, None) is net: 272 continue 273 except: 274 continue 275 if type(getattr(net, member, None)) is PA.TrackedNeuronModule: 276 module = getattr(net, member) 277 if id(module) in seen_ids: 278 continue 279 seen_ids.add(id(module)) 280 this_list = this_list + [module] 281 elif issubclass(type(getattr(net, member, None)), nn.Module): 282 this_list = this_list + get_tracked_modules( 283 getattr(net, member), depth + 1, seen_ids 284 ) 285 return this_list
Get a list of all tracked modules
Parameters
- net (nn.Module): The module to search.
- depth (int): The current depth in the recursion.
Returns
- list: A list of all tracked modules found in the network.
def
get_pai_module_params(net, depth, seen_ids=None):
288def get_pai_module_params(net, depth, seen_ids=None): 289 """Get a list of all neuron module parameters 290 291 Parameters 292 ---------- 293 net : nn.Module 294 The module to search. 295 depth : int 296 The current depth in the recursion. 297 298 Returns 299 ------- 300 list 301 A list of all parameters of neuron modules found in this module. 302 303 """ 304 305 if seen_ids is None: 306 seen_ids = set() 307 all_members = net.__dir__() 308 this_list = [] 309 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 310 for submodule_id, layer in net.named_children(): 311 if isinstance(net.get_submodule(submodule_id), PA.PAINeuronModule): # 312 module = net.get_submodule(submodule_id) 313 if id(module) in seen_ids: 314 continue 315 seen_ids.add(id(module)) 316 for param in module.parameters(): 317 if param.requires_grad: 318 this_list = this_list + [param] 319 else: 320 this_list = this_list + get_pai_module_params( 321 net.get_submodule(submodule_id), depth + 1, seen_ids 322 ) 323 else: 324 for member in all_members: 325 if isinstance(getattr(type(net), member, None), property): 326 continue 327 if getattr(net, member, None) == net: 328 continue 329 if isinstance(getattr(net, member, None), PA.PAINeuronModule): 330 module = getattr(net, member) 331 if id(module) in seen_ids: 332 continue 333 seen_ids.add(id(module)) 334 for param in module.parameters(): 335 if param.requires_grad: 336 this_list = this_list + [param] 337 elif issubclass(type(getattr(net, member, None)), nn.Module): 338 this_list = this_list + get_pai_module_params( 339 getattr(net, member), depth + 1, seen_ids 340 ) 341 return this_list
Get a list of all neuron module parameters
Parameters
- net (nn.Module): The module to search.
- depth (int): The current depth in the recursion.
Returns
- list: A list of all parameters of neuron modules found in this module.
def
get_pai_network_params(net):
344def get_pai_network_params(net): 345 """Get a list of all neuron module parameters 346 347 Parameters 348 ---------- 349 net : nn.Module 350 The full model to search. 351 352 Returns 353 ------- 354 list 355 A list of all parameters of neuron modules found in the network. 356 357 """ 358 param_list = get_pai_module_params(net, 0) 359 return param_list
Get a list of all neuron module parameters
Parameters
- net (nn.Module): The full model to search.
Returns
- list: A list of all parameters of neuron modules found in the network.
def
replace_predefined_modules(start_module):
362def replace_predefined_modules(start_module): 363 """Replace a module with the module from globals list 364 365 Parameters 366 ---------- 367 start_module : nn.Module 368 The module to replace. 369 370 Returns 371 ------- 372 nn.Module 373 The replaced module. 374 375 """ 376 index = GPA.pc.get_modules_to_replace().index(type(start_module)) 377 return GPA.pc.get_replacement_modules()[index](start_module)
Replace a module with the module from globals list
Parameters
- start_module (nn.Module): The module to replace.
Returns
- nn.Module: The replaced module.
def
scan_module_aliases(net):
380def scan_module_aliases(net): 381 """Find alias module paths that point to already-seen module instances. 382 383 Parameters 384 ---------- 385 net : Any PyTorch Module. 386 387 Returns 388 ------- 389 dict[str, str] 390 Mapping of alias module paths to their canonical module paths. 391 """ 392 canonical = {} 393 aliases = {} 394 for name, module in net.named_modules(remove_duplicate=False): 395 if name == "": 396 continue 397 sub_name = "." + name 398 module_id = id(module) 399 if module_id in canonical: 400 aliases[sub_name] = canonical[module_id] 401 else: 402 canonical[module_id] = sub_name 403 return aliases
Find alias module paths that point to already-seen module instances.
Parameters
- net (Any PyTorch Module.):
Returns
- dict[str, str]: Mapping of alias module paths to their canonical module paths.
def
convert_module( net, depth, name_so_far, converted_list, converted_names_list, neuron_module_class, tracked_module_class, config_setup=False):
406def convert_module( 407 net, 408 depth, 409 name_so_far, 410 converted_list, 411 converted_names_list, 412 neuron_module_class, 413 tracked_module_class, 414 config_setup=False, 415): 416 """Recursive function to do all conversion of modules to wrappers of modules 417 418 This is the function that goes through all of the module lists from 419 the globals file and does all the conversion and replacements to 420 setup the dendrite scaffolding as instructed. 421 422 Parameters 423 ---------- 424 net : nn.Module 425 The module to convert. 426 depth : int 427 The current depth in the recursion. 428 name_so_far : str 429 The name of the module so far in the recursion. 430 converted_list : list 431 A list of already converted module ids to avoid infinite loops. 432 converted_names_list : list 433 A corresponding list to help debug duplicate conversions 434 435 Returns 436 ------- 437 nn.Module 438 The converted module. 439 440 """ 441 if GPA.pc.get_verbose(): 442 print("calling convert on %s depth %d" % (net, depth)) 443 print( 444 "calling convert on %s: %s, depth %d" 445 % (name_so_far, type(net).__name__, depth) 446 ) 447 if isinstance(net, neuron_module_class) or ( 448 (tracked_module_class is not None) and isinstance(net, tracked_module_class) 449 ): 450 if GPA.pc.get_verbose(): 451 print( 452 "This is only being called because something in your model " 453 "is pointed to twice by two different variables. Highest " 454 "thing on the list is one of the duplicates" 455 ) 456 return net 457 if depth == 0 and name_so_far == "": 458 aliases = scan_module_aliases(net) 459 existing_not_save = set(GPA.pc.get_module_names_to_not_save()) 460 aliases_to_skip = [ 461 alias for alias in aliases.keys() if alias not in existing_not_save 462 ] 463 if aliases_to_skip and not config_setup: 464 GPA.pc.append_module_names_to_not_save(aliases_to_skip) 465 print( 466 "Auto-detected duplicate module aliases via named_modules; " 467 "keeping first-seen paths and skipping:" 468 ) 469 for alias in aliases_to_skip: 470 print(" - %s (keeps %s)" % (alias, aliases[alias])) 471 all_members = net.__dir__() 472 if GPA.pc.get_extra_verbose(): 473 print("all members:") 474 for member in all_members: 475 print(" - %s" % member) 476 if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList): 477 for submodule_id, layer in net.named_children(): 478 sub_name = name_so_far + "." + str(submodule_id) 479 if sub_name in GPA.pc.get_module_ids_to_track(): 480 if GPA.pc.get_verbose(): 481 print("Seq ID is in track IDs: %s" % sub_name) 482 if tracked_module_class is None: 483 continue 484 if config_setup: 485 net.get_submodule(submodule_id).pai_mode = "tracked" 486 else: 487 setattr( 488 net, 489 submodule_id, 490 tracked_module_class(net.get_submodule(submodule_id), sub_name), 491 ) 492 continue 493 if sub_name in GPA.pc.get_module_ids_to_perforate(): 494 if GPA.pc.get_verbose(): 495 print("Seq ID is in convert IDs: %s" % sub_name) 496 if config_setup: 497 net.get_submodule(submodule_id).pai_mode = "perforated" 498 else: 499 setattr( 500 net, 501 submodule_id, 502 neuron_module_class(net.get_submodule(submodule_id), sub_name), 503 ) 504 continue 505 if type(net.get_submodule(submodule_id)) in GPA.pc.get_modules_to_replace(): 506 if GPA.pc.get_verbose(): 507 print( 508 "Seq sub is in replacement module so replacing: %s" % sub_name 509 ) 510 if config_setup: 511 net.get_submodule(submodule_id).pai_mode = "replaced" 512 else: 513 setattr( 514 net, 515 submodule_id, 516 replace_predefined_modules(net.get_submodule(submodule_id)), 517 ) 518 if ( 519 type(net.get_submodule(submodule_id)) in GPA.pc.get_modules_to_track() 520 ) or ( 521 type(net.get_submodule(submodule_id)).__name__ 522 in GPA.pc.get_module_names_to_track() 523 ): 524 if GPA.pc.get_verbose(): 525 print( 526 "Seq sub is in tracking list so initiating tracked for: %s" 527 % sub_name 528 ) 529 if tracked_module_class is None: 530 continue 531 if config_setup: 532 net.get_submodule(submodule_id).pai_mode = "tracked" 533 else: 534 setattr( 535 net, 536 submodule_id, 537 tracked_module_class(net.get_submodule(submodule_id), sub_name), 538 ) 539 elif ( 540 type(net.get_submodule(submodule_id)) 541 in GPA.pc.get_modules_to_perforate() 542 or type(net.get_submodule(submodule_id)).__name__ 543 in GPA.pc.get_module_names_to_perforate() 544 ): 545 if GPA.pc.get_verbose(): 546 print( 547 "Seq sub is in conversion list so initing PAI for: " 548 "%s" % sub_name 549 ) 550 if ( 551 issubclass( 552 type(net.get_submodule(submodule_id)), 553 torch.nn.modules.batchnorm._BatchNorm, 554 ) 555 or issubclass( 556 type(net.get_submodule(submodule_id)), 557 torch.nn.modules.instancenorm._InstanceNorm, 558 ) 559 or issubclass( 560 type(net.get_submodule(submodule_id)), 561 torch.nn.modules.normalization.LayerNorm, 562 ) 563 ): 564 print( 565 "You have an unwrapped normalization layer, this " 566 "is not recommended: " + name_so_far 567 ) 568 pdb.set_trace() 569 if config_setup: 570 net.get_submodule(submodule_id).pai_mode = "perforated" 571 else: 572 setattr( 573 net, 574 submodule_id, 575 neuron_module_class(net.get_submodule(submodule_id), sub_name), 576 ) 577 else: 578 if net != net.get_submodule(submodule_id): 579 converted_list += [id(net.get_submodule(submodule_id))] 580 converted_names_list += [sub_name] 581 if GPA.pc.get_verbose(): 582 print( 583 "sub is module but in no lists so going deeper: %s" 584 % sub_name 585 ) 586 587 setattr( 588 net, 589 submodule_id, 590 convert_module( 591 net.get_submodule(submodule_id), 592 depth + 1, 593 sub_name, 594 converted_list, 595 converted_names_list, 596 neuron_module_class, 597 tracked_module_class, 598 config_setup, 599 ), 600 ) 601 # else: 602 # print('%s is a self pointer so skipping' % (name_so_far + '[' + str(submodule_id) + ']')) 603 elif type(net) in GPA.pc.get_modules_to_track(): 604 # print('skipping type for returning from call to: %s' % (name_so_far)) 605 return net 606 else: 607 for member in all_members: 608 if isinstance(getattr(type(net), member, None), property): 609 continue 610 # Immediately check if able to get the member, if not skip it 611 try: 612 getattr(net, member, None) 613 except: 614 continue 615 sub_name = name_so_far + "." + member 616 member_obj = getattr(net, member, None) 617 618 if isinstance(member_obj, (torch.nn.Parameter, torch.nn.parameter.Parameter)): 619 if sub_name in GPA.pc.get_parameter_ids_to_track(): 620 if GPA.pc.get_verbose(): 621 print("tracking parameter by ID: %s" % sub_name) 622 if config_setup: 623 member_obj.pai_mode = "tracked" 624 else: 625 member_obj.parameter_type = "neuron" 626 member_obj.wrapped = True 627 continue 628 629 # Track module object ids once at this level so duplicate aliases are 630 # caught consistently (including direct children of the root module). 631 if isinstance(member_obj, nn.Module): 632 if id(member_obj) in converted_list: 633 original_sub_name = converted_names_list[ 634 converted_list.index(id(member_obj)) 635 ] 636 print( 637 "The following module has a duplicate pointer within " 638 "your model: %s" % sub_name 639 ) 640 print("Keeping first pointer: %s" % original_sub_name) 641 print("Skipping duplicate pointer: %s" % sub_name) 642 print( 643 "If you prefer to keep %s and skip %s, add %s to module_names_to_not_save before convert." 644 % (sub_name, original_sub_name, original_sub_name) 645 ) 646 if not config_setup: 647 GPA.pc.append_module_names_to_not_save([sub_name]) 648 continue 649 converted_list += [id(member_obj)] 650 converted_names_list += [sub_name] 651 if sub_name in GPA.pc.get_module_ids_to_track(): 652 if GPA.pc.get_verbose(): 653 print("Seq ID is in track IDs: %s" % sub_name) 654 if tracked_module_class is None: 655 continue 656 if config_setup: 657 getattr(net, member).pai_mode = "tracked" 658 else: 659 setattr( 660 net, member, tracked_module_class(getattr(net, member), sub_name) 661 ) 662 continue 663 if sub_name in GPA.pc.get_module_ids_to_perforate(): 664 if GPA.pc.get_verbose(): 665 print("Seq ID is in convert IDs: %s" % sub_name) 666 if config_setup: 667 getattr(net, member).pai_mode = "perforated" 668 else: 669 setattr( 670 net, member, neuron_module_class(getattr(net, member), sub_name) 671 ) 672 continue 673 if id(getattr(net, member, None)) == id(net): 674 if GPA.pc.get_verbose(): 675 print("member sub is a self pointer: %s" % sub_name) 676 continue 677 if sub_name in GPA.pc.get_module_names_to_not_save(): 678 if GPA.pc.get_verbose(): 679 print("Skipping %s during convert" % sub_name) 680 else: 681 if sub_name == ".base_model": 682 print( 683 "By default skipping base_model. See " 684 '"Safetensors Errors" section of ' 685 "customization.md to include it." 686 ) 687 continue 688 if type(getattr(net, member, None)) in GPA.pc.get_modules_to_replace(): 689 if GPA.pc.get_verbose(): 690 print("sub is in replacement module so replacing: %s" % sub_name) 691 if config_setup: 692 getattr(net, member).pai_mode = "replaced" 693 else: 694 setattr( 695 net, member, replace_predefined_modules(getattr(net, member, None)) 696 ) 697 if ( 698 type(getattr(net, member, None)) in GPA.pc.get_modules_to_track() 699 or type(getattr(net, member, None)).__name__ 700 in GPA.pc.get_module_names_to_track() 701 or sub_name in GPA.pc.get_module_ids_to_track() 702 ): 703 if GPA.pc.get_verbose(): 704 print( 705 "sub is in tracking list so initiating tracked for: %s" 706 % sub_name 707 ) 708 if tracked_module_class is None: 709 continue 710 if config_setup: 711 getattr(net, member).pai_mode = "tracked" 712 else: 713 setattr( 714 net, member, tracked_module_class(getattr(net, member), sub_name) 715 ) 716 elif ( 717 type(getattr(net, member, None)) in GPA.pc.get_modules_to_perforate() 718 or type(getattr(net, member, None)).__name__ 719 in GPA.pc.get_module_names_to_perforate() 720 or (sub_name in GPA.pc.get_module_ids_to_perforate()) 721 ): 722 if GPA.pc.get_verbose(): 723 print( 724 "sub is in conversion list so initiating PAI for: %s" % sub_name 725 ) 726 if config_setup: 727 getattr(net, member).pai_mode = "perforated" 728 else: 729 setattr( 730 net, 731 member, 732 neuron_module_class(getattr(net, member), sub_name), 733 ) 734 elif ( 735 issubclass(type(getattr(net, member, None)), nn.Module) 736 or issubclass(type(getattr(net, member, None)), nn.Sequential) 737 or issubclass(type(getattr(net, member, None)), nn.ModuleList) 738 ): 739 if net != getattr(net, member): 740 if GPA.pc.get_verbose(): 741 print( 742 "sub is module but in no lists so going deeper: %s" 743 % sub_name 744 ) 745 setattr( 746 net, 747 member, 748 convert_module( 749 getattr(net, member), 750 depth + 1, 751 sub_name, 752 converted_list, 753 converted_names_list, 754 neuron_module_class, 755 tracked_module_class, 756 config_setup, 757 ), 758 ) 759 if ( 760 issubclass( 761 type(getattr(net, member, None)), 762 torch.nn.modules.batchnorm._BatchNorm, 763 ) 764 or issubclass( 765 type(getattr(net, member, None)), 766 torch.nn.modules.instancenorm._InstanceNorm, 767 ) 768 or issubclass( 769 type(getattr(net, member, None)), 770 torch.nn.modules.normalization.LayerNorm, 771 ) 772 ): 773 if not GPA.pc.get_unwrapped_modules_confirmed(): 774 print( 775 "potentially found a norm Layer that " 776 "is not accounted for, this is not recommended: %s" % (sub_name) 777 ) 778 print( 779 "Set GPA.pc.set_unwrapped_modules_confirmed(True) to skip " 780 "this next time only if you are unable to debug" 781 ) 782 print( 783 "inspect your network to " 784 "see what the module type containing this layer is." 785 ) 786 print("Then do one of the following:") 787 print( 788 " - Add the module type to " 789 "GPA.pc.get_module_names_to_perforate() to wrap it entirely" 790 ) 791 print( 792 " - If the norm layer is part of a sequential wrap " 793 "it and the previous layer in a PAISequential" 794 ) 795 print( 796 " - If you do not want to add dendrites to this " 797 "module add the type to GPA.pc.get_module_names_to_track()" 798 ) 799 pdb.set_trace() 800 else: 801 if GPA.pc.get_verbose(): 802 if member[0] != "_" or GPA.pc.get_extra_verbose() is True: 803 print("not calling convert on %s depth %d" % (member, depth)) 804 if GPA.pc.get_verbose(): 805 print("returning from call to: %s" % (name_so_far)) 806 return net
Recursive function to do all conversion of modules to wrappers of modules
This is the function that goes through all of the module lists from the globals file and does all the conversion and replacements to setup the dendrite scaffolding as instructed.
Parameters
- net (nn.Module): The module to convert.
- depth (int): The current depth in the recursion.
- name_so_far (str): The name of the module so far in the recursion.
- converted_list (list): A list of already converted module ids to avoid infinite loops.
- converted_names_list (list): A corresponding list to help debug duplicate conversions
Returns
- nn.Module: The converted module.
def
convert_network(net, layer_name=''):
809def convert_network(net, layer_name=""): 810 """Function that calls convert_module and checks results 811 812 Parameters 813 ---------- 814 net : nn.Module 815 The network to convert. 816 layer_name : str, optional 817 The name of the layer if converting a single layer, by default "" 818 819 Returns 820 ------- 821 nn.Module 822 The converted network. 823 824 """ 825 if GPA.pc.get_perforated_backpropagation(): 826 UPB.initialize_pb() 827 MPB.set_main_parameters(net) 828 if type(net) in GPA.pc.get_modules_to_replace(): 829 net = replace_predefined_modules(net) 830 if (type(net) in GPA.pc.get_modules_to_perforate()) or ( 831 type(net).__name__ in GPA.pc.get_module_names_to_perforate() 832 ): 833 if layer_name == "": 834 print( 835 "converting a single layer without a name, add a " 836 "layer_name param to the call" 837 ) 838 sys.exit(-1) 839 net = PA.PAINeuronModule(net, layer_name) 840 else: 841 net = convert_module( 842 net, 0, "", [], [], PA.PAINeuronModule, PA.TrackedNeuronModule 843 ) 844 if GPA.pai_tracker.member_vars["doing_pai"]: 845 missed_ones = [] 846 tracked_ones = [] 847 for name, param in net.named_parameters(): 848 wrapped = "wrapped" in param.__dir__() 849 if wrapped: 850 if GPA.pc.get_verbose(): 851 print("param %s is now wrapped" % (name)) 852 else: 853 tracked = "tracked" in param.__dir__() 854 if tracked: 855 tracked_ones.append(name) 856 else: 857 missed_ones.append(name) 858 if ( 859 len(missed_ones) != 0 or len(tracked_ones) != 0 860 ) and GPA.pc.get_unwrapped_modules_confirmed() is False: 861 print( 862 "\n------------------------------------------------------------------" 863 ) 864 print( 865 "The following params are not wrapped.\n------------------------------------------------------------------" 866 ) 867 for name in tracked_ones: 868 print("." + name) 869 print( 870 "\n------------------------------------------------------------------" 871 ) 872 print( 873 "The following params are not tracked or wrapped.\n------------------------------------------------------------------" 874 ) 875 for name in missed_ones: 876 print("." + name) 877 print( 878 "\n------------------------------------------------------------------" 879 ) 880 print( 881 "Modules that are not wrapped will not have Dendrites to optimize them" 882 ) 883 print( 884 "Modules modules that are not tracked can cause errors and is NOT recommended" 885 ) 886 print( 887 "Any modules in the second list should be added to module_names_to_track" 888 ) 889 890 print( 891 "Set GPA.pc.set_unwrapped_modules_confirmed(True) to skip this next time" 892 ) 893 print( 894 "Inspect your network and see what the module types of these values are to add them to PGB.module_names_to_perforate" 895 ) 896 # If did miss some then set trace to debug 897 if len(missed_ones) != 0: 898 print( 899 "------------------------------------------------------------------\nType 'c' + enter to continue the run to confirm you do not want them to be refined" 900 ) 901 902 pdb.set_trace() 903 print("confirmed") 904 net.register_buffer("tracker_string", torch.tensor([], dtype=torch.uint8)) 905 return net
Function that calls convert_module and checks results
Parameters
- net (nn.Module): The network to convert.
- layer_name (str, optional): The name of the layer if converting a single layer, by default ""
Returns
- nn.Module: The converted network.
def
string_to_tensor(string):
908def string_to_tensor(string): 909 """Helper function to convert a layer_tracker into a string 910 911 This is required for safetensors saving 912 913 Parameters 914 ---------- 915 string : str 916 The string to convert. 917 918 Returns 919 ------- 920 torch.Tensor 921 The converted tensor. 922 923 """ 924 ords = list(map(ord, string)) 925 ords = torch.tensor(ords, dtype=torch.uint8) 926 return ords
Helper function to convert a layer_tracker into a string
This is required for safetensors saving
Parameters
- string (str): The string to convert.
Returns
- torch.Tensor: The converted tensor.
def
string_from_tensor(string_tensor):
929def string_from_tensor(string_tensor): 930 """Convert a tensor back into a string 931 932 Parameters 933 ---------- 934 string_tensor : torch.Tensor 935 The tensor to convert. 936 937 Returns 938 ------- 939 str 940 The converted string. 941 942 """ 943 ords = string_tensor.tolist() 944 to_return = "" 945 # Doing block processing like this helps with memory errors 946 while len(ords) != 0: 947 remaining_ords = ords[100000:] 948 ords = ords[:100000] 949 to_append = "".join(map(chr, ords)) 950 to_return = to_return + to_append 951 ords = remaining_ords 952 return to_return
Convert a tensor back into a string
Parameters
- string_tensor (torch.Tensor): The tensor to convert.
Returns
- str: The converted string.
def
save_system(net, folder, name):
955def save_system(net, folder, name): 956 """Save the entire system 957 958 This saves the network itself as well as the tracker information 959 960 Parameters 961 ---------- 962 net : nn.Module 963 The network to save. 964 folder : str 965 The folder to save the network in. 966 name : str 967 The name to save the network under. 968 969 Returns 970 ------- 971 None 972 973 """ 974 if GPA.pc.get_verbose(): 975 print("saving system %s" % name) 976 temp = string_to_tensor(GPA.pai_tracker.to_string()) 977 if hasattr(net, "tracker_string"): 978 net.tracker_string = string_to_tensor(GPA.pai_tracker.to_string()).to( 979 next(net.parameters()).device 980 ) 981 else: 982 net.register_buffer( 983 "tracker_string", 984 string_to_tensor(GPA.pai_tracker.to_string()).to( 985 next(net.parameters()).device 986 ), 987 ) 988 # Before saving the tracker must be cleared to not contain pointers to the 989 # models modules 990 old_list = GPA.pai_tracker.neuron_module_vector 991 GPA.pai_tracker.neuron_module_vector = [] 992 save_net(net, folder, name) 993 GPA.pai_tracker.neuron_module_vector = old_list 994 pai_save_system(net, folder, name)
Save the entire system
This saves the network itself as well as the tracker information
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
def
load_system( net, folder, name, load_from_restart=False, switch_call=False, load_from_manual_save=False):
997def load_system( 998 net, 999 folder, 1000 name, 1001 load_from_restart=False, 1002 switch_call=False, 1003 load_from_manual_save=False, 1004): 1005 """Load the entire system 1006 1007 This is what should be used to load a saved system and restart training 1008 1009 Parameters 1010 ---------- 1011 net : nn.Module 1012 The network to load into. 1013 folder : str 1014 The folder to load the network from. 1015 name : str 1016 The name to load the network from. 1017 load_from_restart : bool, optional 1018 Whether this is being loaded from an automatic restart, by default False 1019 switch_call : bool, optional 1020 Whether this is being called from a switch, by default False 1021 load_from_manual_save : bool, optional 1022 Whether this is being loaded from a manual save, by default False 1023 1024 Returns 1025 ------- 1026 nn.Module 1027 The loaded network. 1028 1029 Notes 1030 ----- 1031 If you manually call save_system then load_from_manual_save should be True 1032 1033 """ 1034 if GPA.pc.get_verbose(): 1035 print("loading system %s" % name) 1036 net = load_net(net, folder, name) 1037 GPA.pai_tracker.reset_module_vector(net, load_from_restart) 1038 1039 GPA.pai_tracker.from_string(string_from_tensor(net.tracker_string)) 1040 GPA.pai_tracker.saved_time = time.time() 1041 GPA.pai_tracker.loaded = True 1042 GPA.pai_tracker.member_vars["current_best_validation_score"] = 0 1043 GPA.pai_tracker.member_vars["epoch_last_improved"] = GPA.pai_tracker.member_vars[ 1044 "num_epochs_run" 1045 ] 1046 if GPA.pc.get_verbose(): 1047 print( 1048 "after loading epoch last improved is %d mode is %c" 1049 % ( 1050 GPA.pai_tracker.member_vars["epoch_last_improved"], 1051 GPA.pai_tracker.member_vars["mode"], 1052 ) 1053 ) 1054 1055 # Saves always take place before the call to start_epoch so call it here 1056 # when loading to correct off by 1 problems 1057 if (not switch_call) and (not load_from_manual_save): 1058 GPA.pai_tracker.start_epoch(internal_call=True) 1059 return net
Load the entire system
This is what should be used to load a saved system and restart training
Parameters
- net (nn.Module): The network to load into.
- folder (str): The folder to load the network from.
- name (str): The name to load the network from.
- load_from_restart (bool, optional): Whether this is being loaded from an automatic restart, by default False
- switch_call (bool, optional): Whether this is being called from a switch, by default False
- load_from_manual_save (bool, optional): Whether this is being loaded from a manual save, by default False
Returns
- nn.Module: The loaded network.
Notes
If you manually call save_system then load_from_manual_save should be True
def
load_pretrained_model(net, folder, name, remove_dendrite_scaffolding=False):
1062def load_pretrained_model( 1063 net, 1064 folder, 1065 name, 1066 remove_dendrite_scaffolding=False, 1067): 1068 """Load a pretrained perforated model and reset tracker for fresh training. 1069 1070 This function loads a pretrained model's weights and dendrite structure while 1071 resetting all tracker state (epochs, switch history, etc.) to start training 1072 from scratch on a new task. This is useful for transfer learning where you want 1073 pretrained weights but need fresh training dynamics. 1074 1075 Parameters 1076 ---------- 1077 net : nn.Module 1078 The network to load into. 1079 folder : str 1080 The folder containing the pretrained model. 1081 name : str 1082 The name of the checkpoint to load (e.g., 'best_model', 'beforeSwitch_0'). 1083 remove_dendrite_scaffolding : bool, optional 1084 If True, removes dendrite scaffolding for inference or finetuning without 1085 adding more dendrites using blockwise_network and refresh_net. Default False. 1086 1087 Returns 1088 ------- 1089 nn.Module 1090 The loaded network with reset tracker state. 1091 1092 Examples 1093 -------- 1094 Load pretrained weights for continued dendrite training: 1095 >>> model = load_pretrained_model(model, "pretrained-prefc", "beforeSwitch_0") 1096 1097 Load pretrained weights for finetuning without adding more dendrites: 1098 >>> model = load_pretrained_model(model, "pretrained-prefc", "best_model", 1099 ... remove_dendrite_scaffolding=True) 1100 1101 Notes 1102 ----- 1103 This function: 1104 - Loads model weights and dendrite structure from checkpoint 1105 - Resets all epoch counters to -1 (will become 0 after first start_epoch) 1106 - Resets switch history and validation score tracking 1107 - Clears accuracy/loss history arrays 1108 - Optionally removes dendrite scaffolding (no more dendrite additions) 1109 1110 The tracker is reset to behave as if starting fresh training, while keeping 1111 the learned weights and dendrite structure from the pretrained model. 1112 """ 1113 from perforatedai import globals_perforatedai as GPA 1114 1115 if GPA.pc.get_verbose(): 1116 print(f"Loading pretrained model from {folder}/{name}") 1117 1118 # Load the model weights and dendrite structure 1119 net = load_system(net, folder, name, load_from_manual_save=True) 1120 1121 if GPA.pc.get_verbose(): 1122 print("Resetting tracker state for fresh training...") 1123 1124 # Reset structural training state to true initial values. 1125 # Keeping pretrained architecture/weights while zeroing cycle counters avoids 1126 # stale dendrite bookkeeping referencing empty score buffers. 1127 GPA.pai_tracker.reset_module_vector(net, load_from_restart=True) 1128 GPA.pai_tracker.member_vars["mode"] = "n" 1129 GPA.pai_tracker.member_vars["num_dendrites_added"] = 0 1130 GPA.pai_tracker.member_vars["num_dendrites_integrated"] = 0 1131 GPA.pai_tracker.member_vars["num_cycles"] = 0 1132 GPA.pai_tracker.member_vars["num_dendrite_tries"] = 0 1133 GPA.pai_tracker.member_vars["current_n_set_global_best"] = True 1134 1135 # Reset epoch counters 1136 GPA.pai_tracker.member_vars["num_epochs_run"] = -1 1137 GPA.pai_tracker.member_vars["total_epochs_run"] = -1 1138 GPA.pai_tracker.member_vars["epoch_last_improved"] = 0 1139 GPA.pai_tracker.member_vars["last_switch"] = 0 1140 GPA.pai_tracker.member_vars["manual_train_switch"] = False 1141 1142 # Reset switch history 1143 GPA.pai_tracker.member_vars["switch_epochs"] = [] 1144 GPA.pai_tracker.member_vars["n_switch_epochs"] = [] 1145 GPA.pai_tracker.member_vars["p_switch_epochs"] = [] 1146 GPA.pai_tracker.member_vars["param_counts"] = [] 1147 1148 # Reset validation scores and tracking 1149 GPA.pai_tracker.member_vars["current_best_validation_score"] = 0 1150 GPA.pai_tracker.member_vars["global_best_validation_score"] = 0 1151 GPA.pai_tracker.member_vars["running_accuracy"] = 0 1152 1153 # Clear accuracy/loss history arrays 1154 GPA.pai_tracker.member_vars["accuracies"] = [] 1155 GPA.pai_tracker.member_vars["last_improved_accuracies"] = [] 1156 GPA.pai_tracker.member_vars["test_accuracies"] = [] 1157 GPA.pai_tracker.member_vars["n_accuracies"] = [] 1158 GPA.pai_tracker.member_vars["p_accuracies"] = [] 1159 GPA.pai_tracker.member_vars["running_accuracies"] = [] 1160 GPA.pai_tracker.member_vars["training_loss"] = [] 1161 GPA.pai_tracker.member_vars["training_learning_rates"] = [] 1162 GPA.pai_tracker.member_vars["test_scores"] = [] 1163 1164 # Clear extra scores 1165 GPA.pai_tracker.member_vars["extra_scores"] = {} 1166 GPA.pai_tracker.member_vars["extra_scores_without_graphing"] = {} 1167 GPA.pai_tracker.member_vars["n_extra_scores"] = {} 1168 1169 # Keep per-layer dendrite score buffers initialized by reset_module_vector. 1170 1171 # Clear timing arrays 1172 GPA.pai_tracker.member_vars["n_epoch_times"] = [] 1173 GPA.pai_tracker.member_vars["p_epoch_times"] = [] 1174 GPA.pai_tracker.member_vars["n_train_times"] = [] 1175 GPA.pai_tracker.member_vars["p_train_times"] = [] 1176 GPA.pai_tracker.member_vars["n_val_times"] = [] 1177 GPA.pai_tracker.member_vars["p_val_times"] = [] 1178 1179 # Clear overwritten tracking 1180 GPA.pai_tracker.member_vars["overwritten_extras"] = [] 1181 GPA.pai_tracker.member_vars["overwritten_vals"] = [] 1182 GPA.pai_tracker.member_vars["overwritten_epochs"] = 0 1183 1184 # Reset learning rate search state 1185 GPA.pai_tracker.member_vars["initial_lr_test_epoch_count"] = -1 1186 GPA.pai_tracker.member_vars["current_n_learning_rate_initial_skip_steps"] = 0 1187 GPA.pai_tracker.member_vars["last_max_learning_rate_steps"] = 0 1188 GPA.pai_tracker.member_vars["last_max_learning_rate_value"] = -1 1189 GPA.pai_tracker.member_vars["current_cycle_lr_max_scores"] = [] 1190 GPA.pai_tracker.member_vars["current_step_count"] = 0 1191 GPA.pai_tracker.member_vars["committed_to_initial_rate"] = True 1192 GPA.pai_tracker.member_vars["best_mean_score_improved_this_epoch"] = 0 1193 GPA.pai_tracker.member_vars["step_status"] = TPA.STEP_CLEARED 1194 1195 # Reset saved time 1196 GPA.pai_tracker.start_time = time.time() 1197 GPA.pai_tracker.saved_time = 0 1198 1199 # Match tracker initialization behavior so first validation uses epoch 0. 1200 GPA.pai_tracker.start_epoch(internal_call=True) 1201 1202 if GPA.pc.get_verbose(): 1203 print( 1204 f"Tracker reset complete. Dendrites: {GPA.pai_tracker.member_vars['num_dendrites_integrated']}, " 1205 f"Mode: {GPA.pai_tracker.member_vars['mode']}" 1206 ) 1207 1208 # Optionally remove dendrite scaffolding 1209 if remove_dendrite_scaffolding: 1210 if GPA.pc.get_verbose(): 1211 print("Removing dendrite scaffolding (no dendrite additions)...") 1212 1213 from perforatedai import blockwise_perforatedai as BPA 1214 from perforatedai import clean_perforatedai as CPA 1215 1216 net = BPA.blockwise_network(net) 1217 net = CPA.refresh_net(net) 1218 1219 if GPA.pc.get_verbose(): 1220 print("Dendrite scaffolding removed. Model ready for inference or finetuning.") 1221 1222 return net
Load a pretrained perforated model and reset tracker for fresh training.
This function loads a pretrained model's weights and dendrite structure while resetting all tracker state (epochs, switch history, etc.) to start training from scratch on a new task. This is useful for transfer learning where you want pretrained weights but need fresh training dynamics.
Parameters
- net (nn.Module): The network to load into.
- folder (str): The folder containing the pretrained model.
- name (str): The name of the checkpoint to load (e.g., 'best_model', 'beforeSwitch_0').
- remove_dendrite_scaffolding (bool, optional): If True, removes dendrite scaffolding for inference or finetuning without adding more dendrites using blockwise_network and refresh_net. Default False.
Returns
- nn.Module: The loaded network with reset tracker state.
Examples
Load pretrained weights for continued dendrite training:
>>> model = load_pretrained_model(model, "pretrained-prefc", "beforeSwitch_0")
Load pretrained weights for finetuning without adding more dendrites:
>>> model = load_pretrained_model(model, "pretrained-prefc", "best_model",
... remove_dendrite_scaffolding=True)
Notes
This function:
- Loads model weights and dendrite structure from checkpoint
- Resets all epoch counters to -1 (will become 0 after first start_epoch)
- Resets switch history and validation score tracking
- Clears accuracy/loss history arrays
- Optionally removes dendrite scaffolding (no more dendrite additions)
The tracker is reset to behave as if starting fresh training, while keeping the learned weights and dendrite structure from the pretrained model.
def
save_model_with_weight_tying(model, filepath):
1231def save_model_with_weight_tying(model, filepath): 1232 """Save model with safetensors while handling weight tying automatically 1233 1234 Parameters 1235 ---------- 1236 model : Any PyTorch Module. 1237 filepath : Path to filename. 1238 1239 Returns 1240 ------- 1241 dict[str, str] 1242 Mapping from secondary parameter keys to their primary tied key. 1243 """ 1244 state_dict = model.state_dict() 1245 1246 # Find all weight tied parameters 1247 tensor_to_keys = defaultdict(list) 1248 for key, tensor in state_dict.items(): 1249 # Use tensor data pointer as unique identifier 1250 tensor_id = tensor.data_ptr() 1251 tensor_to_keys[tensor_id].append(key) 1252 1253 # Find tied weights (tensors referenced by multiple keys) 1254 tied_weights = {} 1255 keys_to_remove = set() 1256 for tensor_id, keys in tensor_to_keys.items(): 1257 if len(keys) > 1 and not tensor_id == 0: 1258 # Multiple keys reference the same tensor - this is weight tying 1259 # Sort keys for deterministic ordering 1260 keys = sorted(keys) 1261 primary_key = keys[0] # Keep the first key 1262 for secondary_key in keys[1:]: 1263 tied_weights[secondary_key] = primary_key 1264 keys_to_remove.add(secondary_key) 1265 1266 # Remove tied weights from state_dict (keep only primary references) 1267 filtered_state_dict = { 1268 k: v for k, v in state_dict.items() if k not in keys_to_remove 1269 } 1270 1271 # Create metadata for weight tying information 1272 metadata = {} 1273 if tied_weights: 1274 # Store weight tying info as JSON string in metadata 1275 metadata["weight_tying"] = json.dumps(tied_weights) 1276 save_file(filtered_state_dict, filepath, metadata=metadata) 1277 print(f"Saved model with {len(tied_weights)} weight tying relationships") 1278 return tied_weights
Save model with safetensors while handling weight tying automatically
Parameters
model (Any PyTorch Module.):
filepath (Path to filename.):
Returns
- dict[str, str]: Mapping from secondary parameter keys to their primary tied key.
def
load_model_with_weight_tying(model, filepath):
1281def load_model_with_weight_tying(model, filepath): 1282 """Load model from safetensors while restoring weight tying 1283 1284 Parameters 1285 ---------- 1286 model : Any PyTorch module. 1287 filepath : Path to model file. 1288 1289 Returns 1290 ------- 1291 Loaded model. 1292 """ 1293 with safe_open(filepath, framework="pt") as f: 1294 metadata = f.metadata() 1295 state_dict = {key: f.get_tensor(key) for key in f.keys()} 1296 1297 # Restore weight tying if metadata exists 1298 tied_weights = {} 1299 if metadata and "weight_tying" in metadata: 1300 tied_weights = json.loads(metadata["weight_tying"]) 1301 for secondary_key, primary_key in tied_weights.items(): 1302 if primary_key in state_dict: 1303 # Restore the tied reference 1304 state_dict[secondary_key] = state_dict[primary_key] 1305 print(f"Restored weight tying: {secondary_key} -> {primary_key}") 1306 1307 # Handle tracker_string loading with flexible key matching 1308 tracker_key = None 1309 if "tracker_string" in state_dict: 1310 tracker_key = "tracker_string" 1311 else: 1312 # Search for keys containing "tracker_string" 1313 tracker_keys = [key for key in state_dict.keys() if "tracker_string" in key] 1314 if len(tracker_keys) == 1: 1315 tracker_key = tracker_keys[0] 1316 elif len(tracker_keys) > 1: 1317 print(f"Error: Multiple tracker_string keys found: {tracker_keys}") 1318 pdb.set_trace() 1319 else: 1320 print("Error: No tracker_string found in state_dict") 1321 1322 if tracker_key is not None and hasattr(model, "tracker_string"): 1323 model.tracker_string = state_dict[tracker_key] 1324 1325 model.load_state_dict(state_dict) 1326 return model
Load model from safetensors while restoring weight tying
Parameters
model (Any PyTorch module.):
filepath (Path to model file.):
Returns
- Loaded model.
def
save_net(net, folder, name):
1329def save_net(net, folder, name): 1330 """Save the network 1331 1332 This is called within save_system after the tracker has been 1333 turned into a single tensor to be saved as a part of the network 1334 1335 Parameters 1336 ---------- 1337 net : nn.Module 1338 The network to save. 1339 folder : str 1340 The folder to save the network in. 1341 name : str 1342 The name to save the network under. 1343 1344 Returns 1345 ------- 1346 None 1347 1348 """ 1349 # If running a DDP only save with first thread 1350 if "RANK" in os.environ: 1351 if int(os.environ["RANK"]) != 0: 1352 return 1353 if not os.path.isdir(folder): 1354 os.makedirs(folder) 1355 save_point = folder + "/" 1356 if not os.path.isdir(save_point): 1357 os.mkdir(save_point) 1358 for param in net.parameters(): 1359 param.data = param.data.contiguous() 1360 if GPA.pc.get_using_safe_tensors(): 1361 if GPA.pc.get_weight_tying_experimental(): 1362 save_model_with_weight_tying(net, save_point + name + ".pt") 1363 else: 1364 # Strip the . so that the naming is the same for everywhere but it works with state_dict naming 1365 not_save = [ns.lstrip('.') for ns in GPA.pc.get_module_names_to_not_save()] 1366 state_dict = {k: v for k, v in net.state_dict().items() 1367 if not any(k.startswith(ns) for ns in not_save)} 1368 save_file(state_dict, save_point + name + ".pt") 1369 else: 1370 torch.save(net, save_point + name + ".pt")
Save the network
This is called within save_system after the tracker has been turned into a single tensor to be saved as a part of the network
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
def
save_pai_net(net, folder, name):
1417def save_pai_net(net, folder, name): 1418 """Save the final pai network 1419 1420 This can be called after training to save the final network 1421 with all scaffolding removed so only the refined weights remain 1422 1423 Parameters 1424 ---------- 1425 net : nn.Module 1426 The network to save. 1427 folder : str 1428 The folder to save the network in. 1429 name : str 1430 The name to save the network under. 1431 1432 Returns 1433 ------- 1434 None 1435 1436 """ 1437 # if running a DDP only save with first thread 1438 if "RANK" in os.environ: 1439 if int(os.environ["RANK"]) != 0: 1440 return 1441 1442 # print('calling save: %s' % name) 1443 # GPA.pai_tracker.archive_layer() 1444 # These deep copys are required or the real model will also have its layers replaced 1445 net = prepare_final_model(net) 1446 if not os.path.isdir(folder): 1447 os.makedirs(folder) 1448 save_point = folder + "/" 1449 if not os.path.isdir(save_point): 1450 os.mkdir(save_point) 1451 1452 if GPA.pc.get_using_safe_tensors(): 1453 if GPA.pc.get_weight_tying_experimental(): 1454 save_model_with_weight_tying(net, save_point + name + "_pai.pt") 1455 else: 1456 save_file(net.state_dict(), save_point + name + "_pai.pt") 1457 else: 1458 torch.save(net, save_point + name + "_pai.pt")
Save the final pai network
This can be called after training to save the final network with all scaffolding removed so only the refined weights remain
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
def
manual_load_state_dict(model, state_dict):
1461def manual_load_state_dict(model, state_dict): 1462 """Load a state dict into a model key-by-key with relaxed checks. 1463 1464 Parameters 1465 ---------- 1466 model : nn.Module 1467 Model receiving parameters. 1468 state_dict : dict 1469 Source parameter dictionary. 1470 1471 Notes 1472 ----- 1473 Keys configured in ``module_names_to_not_save`` are skipped. 1474 1475 Returns 1476 ------- 1477 None 1478 This function does not return a value. 1479 """ 1480 own_state = model.state_dict() 1481 not_save = [ns.lstrip('.') for ns in GPA.pc.get_module_names_to_not_save()] 1482 for name, param in state_dict.items(): 1483 if any(name.startswith(ns) for ns in not_save): 1484 print("skipping loading %s based on module_names_to_not_save" % name) 1485 continue 1486 if name not in own_state: 1487 print(f"Warning: {name} not found in model state_dict") 1488 continue 1489 if isinstance(param, torch.nn.Parameter): 1490 # Backwards compatibility for serialized parameters 1491 param = param.data 1492 try: 1493 own_state[name].copy_(param) 1494 except Exception as e: 1495 print(f"Error loading {name}: {e}") 1496 print("Manual load complete")
Load a state dict into a model key-by-key with relaxed checks.
Parameters
- model (nn.Module): Model receiving parameters.
- state_dict (dict): Source parameter dictionary.
Notes
Keys configured in module_names_to_not_save are skipped.
Returns
- None: This function does not return a value.
def
load_net(net, folder, name):
1499def load_net(net, folder, name): 1500 """load the network 1501 1502 This is called within load_system after the tracker has been 1503 loaded 1504 1505 Parameters 1506 ---------- 1507 net : nn.Module 1508 The network to save. 1509 folder : str 1510 The folder to save the network in. 1511 name : str 1512 The name to save the network under. 1513 1514 Returns 1515 ------- 1516 nn.Module 1517 The loaded network. 1518 1519 """ 1520 save_point = folder + "/" 1521 if GPA.pc.get_using_safe_tensors(): 1522 model_path = save_point + name + ".pt" 1523 if GPA.pc.get_weight_tying_experimental(): 1524 return load_model_with_weight_tying(net, model_path) 1525 else: 1526 try: 1527 with safe_open(model_path, framework="pt") as f: 1528 metadata = f.metadata() 1529 if metadata and "weight_tying" in metadata: 1530 return load_model_with_weight_tying(net, model_path) 1531 except Exception: 1532 pass 1533 state_dict = load_file(model_path) 1534 else: 1535 # Try three different torch.load call signatures for compatibility across 1536 # different PyTorch versions. All three are attempted before giving up so 1537 # the user can see every error at once and decide how to fix the file. 1538 error_weights_only_false = None 1539 error_no_weights_only = None 1540 error_no_state_dict = None 1541 1542 try: 1543 state_dict = torch.load( 1544 save_point + name + ".pt", 1545 map_location=torch.device("cpu"), 1546 weights_only=False, 1547 ).state_dict() 1548 except Exception as caught_error: 1549 error_weights_only_false = caught_error 1550 1551 if error_weights_only_false is not None: 1552 try: 1553 state_dict = torch.load( 1554 save_point + name + ".pt", 1555 map_location=torch.device("cpu"), 1556 ).state_dict() 1557 except Exception as caught_error: 1558 error_no_weights_only = caught_error 1559 1560 if error_no_weights_only is not None: 1561 try: 1562 state_dict = torch.load( 1563 save_point + name + ".pt", 1564 map_location=torch.device("cpu"), 1565 ) 1566 except Exception as caught_error: 1567 error_no_state_dict = caught_error 1568 1569 if error_no_state_dict is not None: 1570 separator = "\n" + "=" * 60 + "\n" 1571 print(separator.join([ 1572 "\nAll four load attempts failed for: " + save_point + name + ".pt", 1573 "Attempt 1 (weights_only=False, .state_dict()):\n" + str(error_weights_only_false), 1574 "Attempt 2 (default weights_only, .state_dict()):\n" + str(error_no_weights_only), 1575 "Attempt 3 (default weights_only, no .state_dict()):\n" + str(error_no_state_dict), 1576 "You must find a way to fix at least one of the above errors to load successfully.\n", 1577 "Suggested problems to look into:\n", 1578 "1 - Modules which have member variables that are methods, which cant be pickled. See __setstate__ and __getstate__ functions to append." 1579 ])) 1580 raise error_no_state_dict 1581 return load_net_from_dict(net, state_dict)
load the network
This is called within load_system after the tracker has been loaded
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- nn.Module: The loaded network.
def
get_module_base_name(module):
1583def get_module_base_name(module): 1584 """Normalize a wrapped module name for state-dict key lookup. 1585 1586 Parameters 1587 ---------- 1588 module : nn.Module 1589 Module containing a ``name`` attribute. 1590 1591 Returns 1592 ------- 1593 str 1594 Base name with leading dot and optional ``module.`` prefix removed. 1595 """ 1596 module_name = module.name 1597 # This should always be true 1598 if module_name[0] == ".": 1599 # strip "." 1600 module_name = module_name[1:] 1601 # If it was a dataparallel it will also have a module at the start 1602 # so strip that for loading 1603 if module_name[:6] == "module": 1604 module_name = module_name[7:] 1605 return module_name
Normalize a wrapped module name for state-dict key lookup.
Parameters
- module (nn.Module):
Module containing a
nameattribute.
Returns
- str: Base name with leading dot and optional
module.prefix removed.
def
load_net_from_dict(net, state_dict):
1608def load_net_from_dict(net, state_dict): 1609 """load the network 1610 1611 This is called within load_net 1612 1613 Parameters 1614 ---------- 1615 net : nn.Module 1616 The network to save. 1617 state_dict : dict 1618 The state dictionary to load. 1619 1620 Returns 1621 ------- 1622 nn.Module 1623 The loaded network. 1624 1625 """ 1626 if GPA.pc.get_verbose(): 1627 print("loading net from dict") 1628 pai_modules = get_pai_modules(net, 0) 1629 if pai_modules == []: 1630 print( 1631 "PAI load_net and load_system uses a state_dict so it must be\n" 1632 "called with a net after perforate_model has been called" 1633 ) 1634 print( 1635 "This is being flagged because you are attempting to load a model\n" 1636 "that does not have any pai_modules in it. Confirm that you are calling\n" 1637 "perforate_model on the correct model, and the same model is the one\n" 1638 "being passed into add_validation_score" 1639 ) 1640 import pdb # This needs to be here for cython for some reason. 1641 pdb.set_trace() 1642 sys.exit(-1) 1643 if GPA.pc.get_verbose(): 1644 print( 1645 "setting up arrays and simulating cycles for %d pai modules" 1646 % len(pai_modules) 1647 ) 1648 not_save = GPA.pc.get_module_names_to_not_save() 1649 for module in pai_modules: 1650 if any(module.name.startswith(ns) for ns in not_save): 1651 print("skipping loading %s based on module_names_to_not_save" % module.name) 1652 continue 1653 # Set up name to be what will be saved in the state dict 1654 module_name = get_module_base_name(module) 1655 module.clear_dendrites() 1656 for tracker in module.dendrite_module.dendrite_values: 1657 try: 1658 tracker.setup_arrays( 1659 state_dict[ 1660 module_name + ".dendrite_module.dendrite_values.0.dendrite_storage_shape" 1661 ].tolist() 1662 ) 1663 except Exception as e: 1664 print(e) 1665 first_key = next(iter(state_dict.keys())) 1666 error_msg = f""" 1667This value is missing from the state dict. 1668When missing this value it typically means you converted a module but didn't 1669actually use it in your forward and backward pass. 1670 1671Module was: {module.name} 1672 1673There are many reasons this can happen: 1674 16751 - Check your model definition and forward function and ensure this module is 1676 being used properly. 1677 1678 With GPA.pc.set_verbose(True) you can confirm this is the case if you do not 1679 see a "setting d shape for" this module at the first training batch. 1680 1681 If this is the case, and it is correct to not be passing data through it, 1682 set it to be a tracked module with: 1683 GPA.pc.append_module_ids_to_track(["{module.name}"]) to leave it out 1684 16852 - This can happen if you adjusted your model definition after calling 1686 perforate_model (for example with torch.compile). If the module name printed 1687 above does not contain all modules leading to the main definition, this is 1688 likely the case for your problem. Fix by calling perforate_model after all 1689 other model initialization steps. 1690 16913 - This can happen if the model where you called perforate_model and the model 1692 within add_validation_score are not the same. Check if the module above and 1693 .{first_key} have the same prefix. If one starts with .model or .base etc 1694 and the other does not, this is the problem. 1695 16964 - If you are using this module but then not actually including the correct 1697 output tensor in the forward. For example if you are using an LSTM and 1698 forwarding hidden instead of output but your processors are set up to work 1699 with output. 1700 17015 - If you are not properly calling backward at all. If this is the first module 1702 in your network it is more likely this is the problem. One check in these 1703 cases is to make sure you do not call an initial validation score before the 1704 first backward call. If you do this, while testing_dendrite_capacity is True, 1705 this error will be triggered. 1706 17076 - You have converted a module that is in a frozen part of the network and thus 1708 no gradients are flowing. 1709 17107 - You are running multiple experiments at once with the same save_name. When 1711 running concurrent trials be sure to add save_name=<unique_name> to 1712 perforate_model. 1713""" 1714 print(error_msg) 1715 import pdb # This needs to be here for cython for some reason. 1716 pdb.set_trace() 1717 1718 # Perform as many cycles as the state dict has 1719 num_cycles = int(state_dict[module_name + ".dendrite_module.num_cycles"].item()) 1720 if num_cycles > 0: 1721 simulate_cycles(module, num_cycles, doing_pai=True) 1722 # Handle tracker_string loading with flexible key matching 1723 tracker_key = None 1724 if "tracker_string" in state_dict: 1725 tracker_key = "tracker_string" 1726 else: 1727 # Search for keys containing "tracker_string" 1728 tracker_keys = [key for key in state_dict.keys() if "tracker_string" in key] 1729 if len(tracker_keys) == 1: 1730 tracker_key = tracker_keys[0] 1731 elif len(tracker_keys) > 1: 1732 print(f"Error: Multiple tracker_string keys found: {tracker_keys}") 1733 import pdb # This needs to be here for cython for some reason. 1734 pdb.set_trace() 1735 else: 1736 print("Error: No tracker_string found in state_dict") 1737 import pdb # This needs to be here for cython for some reason. 1738 pdb.set_trace() 1739 1740 if hasattr(net, "tracker_string"): 1741 net.tracker_string = state_dict[tracker_key] 1742 else: 1743 net.register_buffer("tracker_string", state_dict[tracker_key]) 1744 try: 1745 load_result = net.load_state_dict(state_dict, strict=False) 1746 not_save_state_names = [ns.lstrip('.') for ns in not_save] 1747 1748 def is_ignored_key(key): 1749 """Check whether a state-dict key should be ignored. 1750 1751 Parameters 1752 ---------- 1753 key : str 1754 State-dict key to test. 1755 1756 Returns 1757 ------- 1758 bool 1759 ``True`` when key belongs to a not-saved namespace. 1760 """ 1761 return any(key.startswith(ns) for ns in not_save_state_names) 1762 1763 missing_keys = [key for key in load_result.missing_keys if not is_ignored_key(key)] 1764 unexpected_keys = [key for key in load_result.unexpected_keys if not is_ignored_key(key)] 1765 1766 if GPA.pc.get_strict_loading() and (missing_keys or unexpected_keys): 1767 raise RuntimeError( 1768 "Error(s) in loading state_dict for %s:\n\tMissing key(s) in state_dict: %s. \n\tUnexpected key(s) in state_dict: %s." 1769 % (type(net).__name__, missing_keys, unexpected_keys) 1770 ) 1771 except Exception as e: 1772 """ 1773 When modules have high depth to them (i.e. modules within modules not number of layers) 1774 PyTorch can have trouble loading state dicts even when they are correct. 1775 This is a workaround to manually load the state dict if this happens. 1776 """ 1777 filtered_net_keys = { 1778 key 1779 for key in net.state_dict().keys() 1780 if not any(key.startswith(ns.lstrip('.')) for ns in not_save) 1781 } 1782 if filtered_net_keys == set(state_dict.keys()): 1783 print("Attempting manual loading of state_dict") 1784 manual_load_state_dict(net, state_dict) 1785 else: 1786 print(f"Error loading state_dict: {e}") 1787 print("If the error is due to missing keys (e.g., from code changes), you can try:") 1788 print(" GPA.pc.set_strict_loading(False)") 1789 print(" Do not change this unless you are certain the missing keys are not important to load and are expected due to code changes or arch changes.") 1790 print("\ntype 'c' to print full state dicts\n") 1791 import pdb # This needs to be here for cython for some reason. 1792 pdb.set_trace() 1793 print("net state dict is:") 1794 print(net.state_dict()) 1795 print("loaded state dict is:") 1796 print(state_dict) 1797 print( 1798 "Try to check differences. Likely is caused by a module not " 1799 "being converted that should be or vice versa" 1800 ) 1801 pdb.set_trace() 1802 net.to(GPA.pc.get_device()) 1803 return net
load the network
This is called within load_net
Parameters
- net (nn.Module): The network to save.
- state_dict (dict): The state dictionary to load.
Returns
- nn.Module: The loaded network.
def
pai_save_system(net, folder, name):
1806def pai_save_system(net, folder, name): 1807 """Save the entire system with scaffolding removed 1808 1809 This is used for the final network for inference after training 1810 1811 Parameters 1812 ---------- 1813 net : nn.Module 1814 The network to save. 1815 folder : str 1816 The folder to save the network in. 1817 name : str 1818 The name to save the network under. 1819 1820 Returns 1821 ------- 1822 None 1823 1824 """ 1825 net.member_vars = {} 1826 for member_var in GPA.pai_tracker.member_vars: 1827 if member_var == "scheduler_instance" or member_var == "optimizer_instance": 1828 continue 1829 net.member_vars[member_var] = GPA.pai_tracker.member_vars[member_var] 1830 pai_save_net(net, folder, name)
Save the entire system with scaffolding removed
This is used for the final network for inference after training
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
def
deep_copy_pai(net):
1833def deep_copy_pai(net): 1834 """Deep copy a PAI network 1835 1836 1837 Parameters 1838 ---------- 1839 net : nn.Module 1840 The network to copy. 1841 1842 Returns 1843 ------- 1844 nn.Module 1845 The copied network. 1846 1847 Notes 1848 ---- 1849 This is required because processors must be cleared before calling copy 1850 1851 """ 1852 # Dont check this stuff if its before the perforate_model has been called and you're just copying a regular model 1853 if(GPA.pai_tracker != []): 1854 # Clear gradients before saving the model 1855 if ((GPA.pai_tracker.member_vars["optimizer_instance"]) is not None) and ( 1856 GPA.pai_tracker.member_vars["optimizer_instance"] != [] 1857 ): 1858 GPA.pai_tracker.member_vars["optimizer_instance"].zero_grad() 1859 GPA.pai_tracker.clear_all_processors() 1860 return copy.deepcopy(net)
Deep copy a PAI network
Parameters
- net (nn.Module): The network to copy.
Returns
- nn.Module: The copied network.
Notes
This is required because processors must be cleared before calling copy
def
prepare_final_model(net):
1863def prepare_final_model(net): 1864 """Prepare model for final save by removing scaffolding. 1865 1866 This performs all cleanup steps to convert a PAI model with scaffolding 1867 into a clean final model ready for inference or distribution. 1868 1869 Parameters 1870 ---------- 1871 net : nn.Module 1872 The network to prepare. 1873 1874 Returns 1875 ------- 1876 nn.Module 1877 The cleaned model with scaffolding removed. 1878 """ 1879 # Deep copy and clean the model (removes scaffolding) 1880 net = deep_copy_pai(net) 1881 net = BPA.blockwise_network(net) 1882 net = deep_copy_pai(net) 1883 net = CL.refresh_net(net) 1884 1885 # Remove tracker_string (not needed for final model) 1886 if hasattr(net, "tracker_string"): 1887 del net.tracker_string 1888 1889 # Make parameters contiguous 1890 for param in net.parameters(): 1891 param.data = param.data.contiguous() 1892 1893 return net
Prepare model for final save by removing scaffolding.
This performs all cleanup steps to convert a PAI model with scaffolding into a clean final model ready for inference or distribution.
Parameters
- net (nn.Module): The network to prepare.
Returns
- nn.Module: The cleaned model with scaffolding removed.
def
pai_save_net(net, folder, name):
1896def pai_save_net(net, folder, name): 1897 """Save the entire system with scaffolding removed 1898 1899 This is called within pai_save_system after the tracker has been 1900 turned into a single tensor to be saved as a part of the network 1901 1902 1903 Parameters 1904 ---------- 1905 net : nn.Module 1906 The network to save. 1907 folder : str 1908 The folder to save the network in. 1909 name : str 1910 The name to save the network under. 1911 1912 Returns 1913 ------- 1914 None 1915 1916 Notes 1917 ---- 1918 For open source implementation this is not as important since 1919 minimal values are already being used. 1920 1921 """ 1922 1923 if GPA.pc.get_perforated_backpropagation(): 1924 UPB.pb_save_net(net, folder, name) 1925 else: 1926 return
Save the entire system with scaffolding removed
This is called within pai_save_system after the tracker has been turned into a single tensor to be saved as a part of the network
Parameters
- net (nn.Module): The network to save.
- folder (str): The folder to save the network in.
- name (str): The name to save the network under.
Returns
- None
Notes
For open source implementation this is not as important since minimal values are already being used.
def
simulate_cycles(module, num_cycles, doing_pai):
1929def simulate_cycles(module, num_cycles, doing_pai): 1930 """Simulate dendrite addition cycles 1931 1932 Simulate the back and forth processes of adding dendrites to build a 1933 pretrained dendrite model before loading weights. Required for loading 1934 dendrite save files from non dendrite initial models. 1935 1936 If calling this by hand to start off a network with dendrites added, 1937 make sure to pass a single datapoint forward and backward through the network 1938 so that array sizes are initialized. 1939 1940 Parameters 1941 ---------- 1942 module : PA.PAINeuronModule 1943 The module to simulate cycles on. 1944 num_cycles : int 1945 The number of cycles to simulate. 1946 doing_pai : bool 1947 Whether to actually do the simulation. 1948 1949 Returns 1950 ------- 1951 None 1952 1953 """ 1954 1955 check_skipped = GPA.pc.get_checked_skipped_modules() 1956 if doing_pai is False: 1957 return 1958 GPA.pc.set_checked_skipped_modules(True) 1959 mode = "n" 1960 for i in range(num_cycles): 1961 if mode == "n": 1962 module.set_mode("p") 1963 module.create_new_dendrite_module() 1964 mode = "p" 1965 else: 1966 module.set_mode("n") 1967 mode = "n" 1968 GPA.pc.set_checked_skipped_modules(check_skipped)
Simulate dendrite addition cycles
Simulate the back and forth processes of adding dendrites to build a pretrained dendrite model before loading weights. Required for loading dendrite save files from non dendrite initial models.
If calling this by hand to start off a network with dendrites added, make sure to pass a single datapoint forward and backward through the network so that array sizes are initialized.
Parameters
- module (PA.PAINeuronModule): The module to simulate cycles on.
- num_cycles (int): The number of cycles to simulate.
- doing_pai (bool): Whether to actually do the simulation.
Returns
- None
def
count_params(net):
1971def count_params(net): 1972 """Count the number of parameters in the network 1973 1974 If doing perforated backpropagation this calls the PB function 1975 which does not count scaffolding parameters since the final model 1976 will not have them. 1977 1978 Parameters 1979 ---------- 1980 net : nn.Module 1981 The network to count parameters in. 1982 1983 Returns 1984 ------- 1985 int 1986 The number of parameters in the network. 1987 1988 """ 1989 if GPA.pc.get_perforated_backpropagation(): 1990 return UPB.pb_count_params(net) 1991 parameters = net.named_parameters() 1992 unique_params = { 1993 p.data_ptr(): p for name, p in parameters if "parent_module" not in name 1994 }.values() 1995 return sum(p.numel() for p in unique_params)
Count the number of parameters in the network
If doing perforated backpropagation this calls the PB function which does not count scaffolding parameters since the final model will not have them.
Parameters
- net (nn.Module): The network to count parameters in.
Returns
- int: The number of parameters in the network.
def
change_learning_modes(net, folder, name, doing_pai):
1998def change_learning_modes(net, folder, name, doing_pai): 1999 """Change between neuron and dendrite learning modes 2000 2001 High level steps for entire system to switch back and forth between 2002 neuron learning and dendrite learning 2003 2004 Parameters 2005 ---------- 2006 net : nn.Module 2007 The network to change modes on. 2008 folder : str 2009 The folder to save/load the network in/from. 2010 name : str 2011 The name to save/load the network under. 2012 doing_pai : bool 2013 Whether to add dendrites when changing modes. 2014 2015 Returns 2016 ------- 2017 int 2018 The number of parameters in the network. 2019 2020 Notes 2021 ----- 2022 If doing_pai is False this just allows training to continue longer rather than early stopping 2023 2024 """ 2025 # If not adding dendrites this just allows training to continue longer with flags 2026 # every time early stopping should be occurring 2027 if doing_pai is False: 2028 GPA.pai_tracker.member_vars["switch_epochs"].append( 2029 GPA.pai_tracker.member_vars["num_epochs_run"] 2030 ) 2031 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2032 "switch_epochs" 2033 ][-1] 2034 GPA.pai_tracker.reset_vals_for_score_reset() 2035 return net 2036 if GPA.pai_tracker.member_vars["mode"] == "n": 2037 current_epoch = GPA.pai_tracker.member_vars["num_epochs_run"] 2038 overwritten_epochs = GPA.pai_tracker.member_vars["overwritten_epochs"] 2039 overwritten_extra = GPA.pai_tracker.member_vars["extra_scores"] 2040 if GPA.pc.get_drawing_pai(): 2041 overwritten_val = GPA.pai_tracker.member_vars["accuracies"] 2042 else: 2043 overwritten_val = GPA.pai_tracker.member_vars["neuron_accuracies"] 2044 """ 2045 If true don't load the best system 2046 because it will delete dendrites if the previous best was better than 2047 the current best 2048 """ 2049 if not GPA.pc.get_silent(): 2050 print("Importing best Model for switch to PA...") 2051 net = load_system(net, folder, name, switch_call=True) 2052 GPA.pai_tracker.set_dendrite_training() 2053 GPA.pai_tracker.member_vars["overwritten_epochs"] = overwritten_epochs 2054 GPA.pai_tracker.member_vars["overwritten_epochs"] += ( 2055 current_epoch - GPA.pai_tracker.member_vars["num_epochs_run"] 2056 ) 2057 GPA.pai_tracker.member_vars["total_epochs_run"] = ( 2058 GPA.pai_tracker.member_vars["num_epochs_run"] 2059 + GPA.pai_tracker.member_vars["overwritten_epochs"] 2060 ) 2061 2062 if GPA.pc.get_save_old_graph_scores(): 2063 GPA.pai_tracker.member_vars["overwritten_extras"].append(overwritten_extra) 2064 GPA.pai_tracker.member_vars["overwritten_vals"].append(overwritten_val) 2065 else: 2066 GPA.pai_tracker.member_vars["overwritten_extras"] = [overwritten_extra] 2067 GPA.pai_tracker.member_vars["overwritten_vals"] = [overwritten_val] 2068 if GPA.pc.get_drawing_pai(): 2069 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2070 GPA.pai_tracker.member_vars["num_epochs_run"] 2071 ) 2072 else: 2073 if len(GPA.pai_tracker.member_vars["switch_epochs"]) == 0: 2074 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2075 GPA.pai_tracker.member_vars["num_epochs_run"] 2076 ) 2077 else: 2078 GPA.pai_tracker.member_vars["n_switch_epochs"].append( 2079 GPA.pai_tracker.member_vars["n_switch_epochs"][-1] 2080 + ( 2081 (GPA.pai_tracker.member_vars["num_epochs_run"]) 2082 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2083 ) 2084 ) 2085 2086 GPA.pai_tracker.member_vars["switch_epochs"].append( 2087 GPA.pai_tracker.member_vars["num_epochs_run"] 2088 ) 2089 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2090 "switch_epochs" 2091 ][-1] 2092 2093 # Because open source version is only doing neuron training for 2094 # gradient descent dendrites, switch back to n mode right away 2095 if ( 2096 not GPA.pc.get_perforated_backpropagation() 2097 ) or GPA.pc.get_no_extra_n_modes(): 2098 net = change_learning_modes(net, folder, name, doing_pai) 2099 else: 2100 if not GPA.pc.get_silent(): 2101 print("Switching back to N...") 2102 set_best = GPA.pai_tracker.member_vars["current_n_set_global_best"] 2103 GPA.pai_tracker.set_neuron_training() 2104 if len(GPA.pai_tracker.member_vars["p_switch_epochs"]) == 0: 2105 GPA.pai_tracker.member_vars["p_switch_epochs"].append( 2106 ( 2107 (GPA.pai_tracker.member_vars["num_epochs_run"] - 1) 2108 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2109 ) 2110 ) 2111 else: 2112 GPA.pai_tracker.member_vars["p_switch_epochs"].append( 2113 GPA.pai_tracker.member_vars["p_switch_epochs"][-1] 2114 + ( 2115 (GPA.pai_tracker.member_vars["num_epochs_run"]) 2116 - (GPA.pai_tracker.member_vars["switch_epochs"][-1]) 2117 ) 2118 ) 2119 GPA.pai_tracker.member_vars["switch_epochs"].append( 2120 GPA.pai_tracker.member_vars["num_epochs_run"] 2121 ) 2122 GPA.pai_tracker.member_vars["last_switch"] = GPA.pai_tracker.member_vars[ 2123 "switch_epochs" 2124 ][-1] 2125 # Will be false for open source implementation 2126 if GPA.pc.get_retain_all_dendrites() or ( 2127 GPA.pc.get_learn_dendrites_live() and set_best 2128 ): 2129 if not GPA.pc.get_silent(): 2130 print( 2131 "Saving model before starting normal training to " 2132 "retain PBNodes regardless of next N Phase results" 2133 ) 2134 save_system(net, folder, name) 2135 # if its just doing P for learn PAI live then switch back immediately 2136 if GPA.pc.get_perforated_backpropagation() and GPA.pc.get_no_extra_n_modes(): 2137 net = change_learning_modes(net, folder, name, doing_pai) 2138 2139 GPA.pai_tracker.member_vars["param_counts"].append(count_params(net)) 2140 2141 return net
Change between neuron and dendrite learning modes
High level steps for entire system to switch back and forth between neuron learning and dendrite learning
Parameters
- net (nn.Module): The network to change modes on.
- folder (str): The folder to save/load the network in/from.
- name (str): The name to save/load the network under.
- doing_pai (bool): Whether to add dendrites when changing modes.
Returns
- int: The number of parameters in the network.
Notes
If doing_pai is False this just allows training to continue longer rather than early stopping
def
find_param_name_by_id(model, param_id):
2144def find_param_name_by_id(model, param_id): 2145 """ 2146 This is only used for debugging. 2147 Return the fully-qualified parameter name (e.g. "layer1.conv.weight") 2148 for the parameter whose id matches param_id. Returns None if not found. 2149 2150 This uses model.named_parameters(), which already recurses through submodules. 2151 2152 Parameters 2153 ---------- 2154 model : Model to look for param id. 2155 param_id : pointer to a parameter. 2156 2157 Returns 2158 ------- 2159 String representing the module within the model or None if not found. 2160 """ 2161 for name, p in model.named_parameters(recurse=True): 2162 if id(p) == param_id: 2163 return "." + name 2164 return None
This is only used for debugging. Return the fully-qualified parameter name (e.g. "layer1.conv.weight") for the parameter whose id matches param_id. Returns None if not found.
This uses model.named_parameters(), which already recurses through submodules.
Parameters
model (Model to look for param id.):
param_id (pointer to a parameter.):
Returns
- String representing the module within the model or None if not found.
def
add_method_delegation_to_module(wrapper_module, method_name):
2167def add_method_delegation_to_module(wrapper_module, method_name): 2168 """Add delegating methods to a wrapper module that has a main_module attribute. 2169 2170 This adds the specified methods to the wrapper module instance so they 2171 properly delegate to the wrapped main_module. Works for any wrapper module 2172 (TrackedNeuronModule, PAINeuronModule, etc.) that has a main_module attribute. 2173 2174 Args: 2175 wrapper_module: A wrapper module instance with a main_module attribute 2176 method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') 2177 2178 Parameters 2179 ---------- 2180 wrapper_module : PyTorch Module that contains a sub module. 2181 method_name : Name of method to wrap. 2182 2183 Returns 2184 ------- 2185 None 2186 """ 2187 import types 2188 2189 if hasattr(wrapper_module.main_module, method_name): 2190 # Create a delegating method that forwards to main_module 2191 def make_delegated_method(name): 2192 """Create a bound delegation function for a given attribute name. 2193 2194 Parameters 2195 ---------- 2196 name : str 2197 Attribute name to forward to ``main_module``. 2198 2199 Returns 2200 ------- 2201 callable 2202 Function that delegates access or invocation. 2203 """ 2204 def delegated_method(self, *args, **kwargs): 2205 """Delegate attribute access or method call to ``main_module``. 2206 2207 Parameters 2208 ---------- 2209 *args : tuple 2210 Positional arguments forwarded to delegated callables. 2211 **kwargs : dict 2212 Keyword arguments forwarded to delegated callables. 2213 2214 Returns 2215 ------- 2216 Any 2217 Delegated attribute value or method result. 2218 """ 2219 main_module_attr = getattr(self.main_module, name, None) 2220 if main_module_attr is None: 2221 raise AttributeError( 2222 f"'{type(self.main_module).__name__}' object has no attribute '{name}'" 2223 ) 2224 if callable(main_module_attr): 2225 return main_module_attr(*args, **kwargs) 2226 return main_module_attr 2227 2228 return delegated_method 2229 2230 # Bind it to this specific instance 2231 setattr( 2232 wrapper_module, 2233 method_name, 2234 types.MethodType(make_delegated_method(method_name), wrapper_module), 2235 )
Add delegating methods to a wrapper module that has a main_module attribute.
This adds the specified methods to the wrapper module instance so they properly delegate to the wrapped main_module. Works for any wrapper module (TrackedNeuronModule, PAINeuronModule, etc.) that has a main_module attribute.
Args: wrapper_module: A wrapper module instance with a main_module attribute method_name: The method name to delegate (e.g., '_gradient_checkpointing_func')
Parameters
wrapper_module (PyTorch Module that contains a sub module.):
method_name (Name of method to wrap.):
Returns
- None
def
apply_method_delegation_to_model(model, method_name, main_module_type):
2238def apply_method_delegation_to_model(model, method_name, main_module_type): 2239 """Recursively apply method delegation to all wrapper modules with main_module in a model. 2240 2241 This traverses the entire model and adds method delegation for any module that has 2242 a main_module attribute and optionally matches specified types. 2243 2244 Args: 2245 model: The PyTorch model to traverse 2246 method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') 2247 main_module_type: main_module type name to filter by. 2248 Example: 'Qwen2DecoderLayer' 2249 2250 Example: 2251 # Apply gradient checkpointing delegation to all decoder layers 2252 apply_method_delegation_to_model( 2253 model, 2254 '_gradient_checkpointing_func', 2255 main_module_type='Qwen2DecoderLayer' 2256 ) 2257 2258 Parameters 2259 ---------- 2260 model : PyTorch model. 2261 method_name : method to delegate. 2262 main_module_type : type of module that has this method. 2263 2264 Returns 2265 ------- 2266 None 2267 This function does not return a value. 2268 """ 2269 count = 0 2270 for name, module in model.named_modules(): 2271 # Check if module has main_module attribute (it's a wrapper) 2272 if hasattr(module, "main_module"): 2273 # Check if we should apply based on main_module type 2274 should_apply = True 2275 if main_module_type is not None: 2276 main_module_type_name = type(module.main_module).__name__ 2277 should_apply = main_module_type_name == main_module_type 2278 2279 if should_apply: 2280 add_method_delegation_to_module(module, method_name) 2281 count += 1 2282 2283 print(f"[PAI] Applied method delegation to {count} wrapper module instances")
Recursively apply method delegation to all wrapper modules with main_module in a model.
This traverses the entire model and adds method delegation for any module that has a main_module attribute and optionally matches specified types.
Args: model: The PyTorch model to traverse method_name: The method name to delegate (e.g., '_gradient_checkpointing_func') main_module_type: main_module type name to filter by. Example: 'Qwen2DecoderLayer'
Example: # Apply gradient checkpointing delegation to all decoder layers apply_method_delegation_to_model( model, '_gradient_checkpointing_func', main_module_type='Qwen2DecoderLayer' )
Parameters
model (PyTorch model.):
method_name (method to delegate.):
main_module_type (type of module that has this method.):
Returns
- None: This function does not return a value.
def
make_json_serializable(obj):
2286def make_json_serializable(obj): 2287 """Recursively convert non-JSON-serializable objects to strings. 2288 2289 Parameters 2290 ---------- 2291 obj : any 2292 The object to convert 2293 2294 Returns 2295 ------- 2296 Any 2297 JSON-serializable version of the object 2298 """ 2299 if isinstance(obj, (str, int, float, bool, type(None))): 2300 return obj 2301 elif isinstance(obj, dict): 2302 return {k: make_json_serializable(v) for k, v in obj.items()} 2303 elif isinstance(obj, (list, tuple)): 2304 return [make_json_serializable(item) for item in obj] 2305 else: 2306 # Convert non-serializable types to string 2307 return str(obj)
Recursively convert non-JSON-serializable objects to strings.
Parameters
- obj (any): The object to convert
Returns
- Any: JSON-serializable version of the object
def
extract_gpa_config():
2310def extract_gpa_config(): 2311 """Extract all configuration from GPA.pc by calling all get_* methods. 2312 2313 Returns 2314 ------- 2315 dict[str, Any] 2316 Dictionary with all GPA.pc configuration values and type metadata 2317 2318 Examples 2319 -------- 2320 >>> config = extract_gpa_config() 2321 >>> # Returns: {'max_dendrites': 10, 'device': 'cuda', '_types': {...}} 2322 2323 Parameters 2324 ---------- 2325 None 2326 2327 """ 2328 config = {} 2329 config_types = {} 2330 2331 # Get all attributes from GPA.pc 2332 for attr_name in dir(GPA.pc): 2333 # Check if it starts with 'get_' 2334 if attr_name.startswith("get_"): 2335 try: 2336 # Get the method 2337 method = getattr(GPA.pc, attr_name) 2338 2339 # Check if it's callable 2340 if callable(method): 2341 # Call it and store result with key as name without 'get_' 2342 key = attr_name[4:] # Remove 'get_' prefix 2343 value = method() 2344 2345 # Check if this is an array (has corresponding append_ method) 2346 append_method_name = f"append_{key}" 2347 is_array = hasattr(GPA.pc, append_method_name) 2348 2349 if is_array and isinstance(value, (list, tuple)): 2350 # Store array element type 2351 if len(value) > 0: 2352 element_type = type(value[0]).__name__ 2353 else: 2354 element_type = None # empty array, no conversion needed 2355 config_types[key] = { 2356 "is_array": True, 2357 "element_type": element_type, 2358 } 2359 else: 2360 # Store value type 2361 config_types[key] = { 2362 "is_array": False, 2363 "type": type(value).__name__, 2364 } 2365 2366 # Make sure value is JSON serializable 2367 config[key] = make_json_serializable(value) 2368 except Exception as e: 2369 # Skip if method fails 2370 if GPA.pc.get_verbose(): 2371 print(f"Skipping {attr_name}: {e}") 2372 continue 2373 2374 # Add types metadata to config 2375 config["_types"] = config_types 2376 2377 return config
Extract all configuration from GPA.pc by calling all get_* methods.
Returns
- dict[str, Any]: Dictionary with all GPA.pc configuration values and type metadata
Examples
>>> config = extract_gpa_config()
>>> # Returns: {'max_dendrites': 10, 'device': 'cuda', '_types': {...}}
Parameters
- None
def
convert_to_type(value, type_name):
2380def convert_to_type(value, type_name): 2381 """Convert a value to the specified type. 2382 2383 Parameters 2384 ---------- 2385 value : any 2386 The value to convert 2387 type_name : str 2388 The target type name 2389 2390 Returns 2391 ------- 2392 Any 2393 The converted value 2394 """ 2395 if type_name == "NoneType" or value is None: 2396 return None 2397 elif type_name == "bool": 2398 if isinstance(value, str): 2399 return value.lower() in ("true", "1", "yes") 2400 return bool(value) 2401 elif type_name == "int": 2402 return int(value) 2403 elif type_name == "float": 2404 return float(value) 2405 elif type_name == "str": 2406 return str(value) 2407 elif type_name == "list": 2408 if not isinstance(value, list): 2409 return [value] 2410 return value 2411 elif type_name == "dict": 2412 if not isinstance(value, dict): 2413 return {} 2414 return value 2415 elif type_name == "type": 2416 # Handle type objects - convert string representation back to type 2417 if isinstance(value, str): 2418 # Try to evaluate the type string (e.g., "<class 'torch.nn.Linear'>") 2419 # Extract the class path from the string 2420 if value.startswith("<class '") and value.endswith("'>"): 2421 class_path = value[ 2422 8:-2 2423 ] # Extract 'torch.nn.Linear' from "<class 'torch.nn.Linear'>" 2424 parts = class_path.split(".") 2425 # Try to import and get the type 2426 try: 2427 module_name = ".".join(parts[:-1]) 2428 class_name = parts[-1] 2429 module = __import__(module_name, fromlist=[class_name]) 2430 return getattr(module, class_name) 2431 except Exception as e: 2432 print( 2433 f"Warning: Could not convert type string '{value}' to actual type: {e}" 2434 ) 2435 return value 2436 return value 2437 return value 2438 elif type_name == "dtype": 2439 # Handle torch dtype objects 2440 if isinstance(value, str): 2441 # Convert string like "torch.float32" to actual dtype 2442 import torch 2443 2444 try: 2445 # Try to get the dtype from torch module 2446 if value.startswith("torch."): 2447 dtype_name = value.split(".")[ 2448 1 2449 ] # Get 'float32' from 'torch.float32' 2450 return getattr(torch, dtype_name) 2451 else: 2452 return getattr(torch, value) 2453 except Exception as e: 2454 print( 2455 f"Warning: Could not convert dtype string '{value}' to actual dtype: {e}" 2456 ) 2457 return value 2458 return value 2459 elif type_name == "device": 2460 # Handle torch device objects 2461 if isinstance(value, str): 2462 # Convert string like "cuda" or "cpu" to torch.device 2463 import torch 2464 2465 try: 2466 return torch.device(value) 2467 except Exception as e: 2468 print( 2469 f"Warning: Could not convert device string '{value}' to actual device: {e}" 2470 ) 2471 return value 2472 return value 2473 elif type_name == "builtin_function_or_method": 2474 # Handle torch functions like torch.sigmoid, torch.relu, etc. 2475 if isinstance(value, str): 2476 # Parse string like "<built-in method sigmoid of type object at 0x...>" 2477 # to extract the function name 2478 import torch 2479 2480 try: 2481 if "<built-in method " in value and " of type object" in value: 2482 # Extract function name between '<built-in method ' and ' of type object' 2483 start = value.find("<built-in method ") + len("<built-in method ") 2484 end = value.find(" of type object") 2485 func_name = value[start:end] 2486 # Try to get the function from torch module 2487 if hasattr(torch, func_name): 2488 return getattr(torch, func_name) 2489 else: 2490 print(f"Warning: torch.{func_name} not found") 2491 return value 2492 else: 2493 return value 2494 except Exception as e: 2495 print( 2496 f"Warning: Could not convert builtin function string '{value}': {e}" 2497 ) 2498 return value 2499 return value 2500 else: 2501 # Unknown type - error and debug 2502 print(f"ERROR: Unknown type '{type_name}' for value: {value}") 2503 print(f"Type of value is: {type(value).__name__}") 2504 pdb.set_trace() 2505 return value
Convert a value to the specified type.
Parameters
- value (any): The value to convert
- type_name (str): The target type name
Returns
- Any: The converted value
def
convert_to_type_array(value, element_type):
2508def convert_to_type_array(value, element_type): 2509 """Convert an array's elements to the specified type. 2510 2511 Parameters 2512 ---------- 2513 value : list or tuple 2514 The array to convert 2515 element_type : str or None 2516 The target type name for elements, None if array was empty 2517 2518 Returns 2519 ------- 2520 list 2521 The array with converted elements 2522 """ 2523 if not isinstance(value, (list, tuple)): 2524 return value 2525 # If element_type is None (empty array), no conversion needed 2526 if element_type is None: 2527 return list(value) if isinstance(value, tuple) else value 2528 return [convert_to_type(item, element_type) for item in value]
Convert an array's elements to the specified type.
Parameters
- value (list or tuple): The array to convert
- element_type (str or None): The target type name for elements, None if array was empty
Returns
- list: The array with converted elements
def
set_gpa_config(config):
2531def set_gpa_config(config): 2532 """Set GPA.pc configuration by calling all set_* methods. 2533 2534 This is the reverse of extract_gpa_config(). It takes a configuration 2535 dictionary and calls the corresponding set_* methods on GPA.pc. 2536 Uses type metadata to ensure values are converted to the correct type. 2537 2538 Parameters 2539 ---------- 2540 config : dict 2541 Dictionary with configuration values (keys without 'set_' prefix) 2542 and optional '_types' metadata 2543 2544 Examples 2545 -------- 2546 >>> config = {'verbose': True, 'device': 'cuda'} 2547 >>> set_gpa_config(config) 2548 # Calls GPA.pc.set_verbose(True), GPA.pc.set_device('cuda'), etc. 2549 2550 Returns 2551 ------- 2552 Count of parameters that were set 2553 """ 2554 set_count = 0 2555 skip_count = 0 2556 2557 # Extract type information 2558 config_types = config.get("_types", {}) 2559 2560 for key, value in config.items(): 2561 # Skip the types metadata 2562 if key == "_types": 2563 continue 2564 2565 # Construct the set method name 2566 set_method_name = f"set_{key}" 2567 2568 # Check if the set method exists 2569 if hasattr(GPA.pc, set_method_name): 2570 try: 2571 method = getattr(GPA.pc, set_method_name) 2572 if callable(method): 2573 # Convert value to correct type if we have type info 2574 if key in config_types: 2575 type_info = config_types[key] 2576 if type_info.get("is_array", False): 2577 # Convert array elements to correct type 2578 element_type = type_info.get("element_type", "str") 2579 value = convert_to_type_array(value, element_type) 2580 else: 2581 # Convert single value to correct type 2582 value_type = type_info.get("type", "str") 2583 value = convert_to_type(value, value_type) 2584 2585 method(value) 2586 set_count += 1 2587 if GPA.pc.get_verbose(): 2588 print(f"Set {key} = {value}") 2589 except Exception as e: 2590 skip_count += 1 2591 if GPA.pc.get_verbose(): 2592 print(f"Failed to set {key}: {e}") 2593 else: 2594 skip_count += 1 2595 if GPA.pc.get_verbose(): 2596 print(f"No setter found for {key} (looking for {set_method_name})") 2597 2598 if GPA.pc.get_verbose(): 2599 print(f"Applied {set_count} PAI configuration settings ({skip_count} skipped)") 2600 2601 return set_count
Set GPA.pc configuration by calling all set_* methods.
This is the reverse of extract_gpa_config(). It takes a configuration dictionary and calls the corresponding set_* methods on GPA.pc. Uses type metadata to ensure values are converted to the correct type.
Parameters
- config (dict): Dictionary with configuration values (keys without 'set_' prefix) and optional '_types' metadata
Examples
>>> config = {'verbose': True, 'device': 'cuda'}
>>> set_gpa_config(config)
<h1 id="calls-gpapcset_verbosetrue-gpapcset_devicecuda-etc">Calls GPA.pc.set_verbose(True), GPA.pc.set_device('cuda'), etc.</h1>
Returns
- Count of parameters that were set
def
upload_to_huggingface(*args, **kwargs):
2811 def upload_to_huggingface(*args, **kwargs): 2812 """Raise an informative error when HuggingFace dependencies are missing. 2813 2814 Parameters 2815 ---------- 2816 *args : tuple[Any, ...] 2817 Positional arguments accepted for API compatibility. 2818 **kwargs : dict[str, Any] 2819 Keyword arguments accepted for API compatibility. 2820 2821 Returns 2822 ------- 2823 None 2824 Always raises ``ImportError``. 2825 """ 2826 raise ImportError( 2827 "huggingface_hub is required for upload_to_huggingface. " 2828 "Install it with: pip install huggingface_hub" 2829 )
Raise an informative error when HuggingFace dependencies are missing.
Parameters
- *args (tuple[Any, ...]): Positional arguments accepted for API compatibility.
- **kwargs (dict[str, Any]): Keyword arguments accepted for API compatibility.
Returns
- None: Always raises
ImportError.
def
from_hf_pretrained(*args, **kwargs):
2831 def from_hf_pretrained(*args, **kwargs): 2832 """Raise an informative error when HuggingFace dependencies are missing. 2833 2834 Parameters 2835 ---------- 2836 *args : tuple[Any, ...] 2837 Positional arguments accepted for API compatibility. 2838 **kwargs : dict[str, Any] 2839 Keyword arguments accepted for API compatibility. 2840 2841 Returns 2842 ------- 2843 None 2844 Always raises ``ImportError``. 2845 """ 2846 raise ImportError( 2847 "huggingface_hub is required for from_hf_pretrained. " 2848 "Install it with: pip install huggingface_hub" 2849 )
Raise an informative error when HuggingFace dependencies are missing.
Parameters
- *args (tuple[Any, ...]): Positional arguments accepted for API compatibility.
- **kwargs (dict[str, Any]): Keyword arguments accepted for API compatibility.
Returns
- None: Always raises
ImportError.