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