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