perforatedai.library_perforatedai
1# Copyright (c) 2025 Perforated AI 2 3import math 4import pdb 5from itertools import chain 6 7import torch 8import torch.nn as nn 9import torch.nn.init as init 10import torch.nn.functional as F 11import torchvision.models.resnet as resnet_pt 12from abc import ABC, abstractmethod 13 14from perforatedai import globals_perforatedai as GPA 15 16""" 17Details on processors can be found in customization.md in the API directory. 18 19They exist to enable simplicity in adding dendrites to modules where 20forward() is not one tensor in and one tensor out. 21 22The main module has one instance, which uses post_n1 and post_n2 23and each new Dendrite node gets a unique instance to use pre_d and post_d. 24""" 25 26 27class PAIProcessor(ABC): 28 """ 29 Abstract base class for processing neuron and dendrite operations. 30 31 Processors handle state management and data flow between neurons and 32 dendrites, allowing for custom pre/post processing of modules which have 33 multiple inputs and outputs, rather than the default single tensor input/output. 34 Subclasses should implement the five core processing methods to handle 35 their specific state management needs. 36 """ 37 38 @abstractmethod 39 def post_n1(self, *args, **kwargs): 40 """ 41 Post-process neuron output before dendrite processing. 42 43 Called immediately after the main module/neuron is executed and before 44 any dendrite processing occurs. This method should extract and return 45 only the tensor of the neuron output that should be seen by 46 dendrite operations. 47 48 Parameters 49 ---------- 50 *args : tuple 51 Positional arguments, typically containing the neuron output. 52 **kwargs : dict 53 Keyword arguments from the neuron output. 54 55 Returns 56 ------- 57 Any 58 The filtered output to be passed to dendrite processing. 59 """ 60 pass 61 62 @abstractmethod 63 def post_n2(self, *args, **kwargs): 64 """ 65 Post-process dendrite-modified output before final return. 66 67 Called after dendrite processing is complete and before passing the 68 final value forward in the network. This method should combine the 69 dendrite-modified output with any stored state to produce the complete 70 output that matches the expected format of the main module. 71 72 Parameters 73 ---------- 74 *args : tuple 75 Positional arguments containing the dendrite-modified output. 76 **kwargs : dict 77 Keyword arguments from the processing chain. 78 79 Returns 80 ------- 81 Any 82 The complete output in the format expected by downstream components. 83 """ 84 pass 85 86 @abstractmethod 87 def pre_d(self, *args, **kwargs): 88 """ 89 Pre-process input before dendrite operations. 90 91 Filters and prepares inputs for dendrite processing. This method handles 92 special cases such as initial time steps vs. subsequent iterations, 93 ensuring dendrites receive the appropriate inputs (e.g., external inputs 94 vs. internal recurrent state). 95 96 Parameters 97 ---------- 98 *args : tuple 99 Positional arguments containing inputs to the PAI module. 100 **kwargs : dict 101 Keyword arguments containing inputs to the PAI module. 102 103 Returns 104 ------- 105 tuple 106 A tuple of (processed_args, processed_kwargs) to pass to dendrite. 107 """ 108 pass 109 110 @abstractmethod 111 def post_d(self, *args, **kwargs): 112 """ 113 Post-process dendrite output and manage state. 114 115 Processes the output from dendrite operations, storing any state needed 116 for future iterations and returning only the portion that should be 117 combined with the neuron output. E.g. this is where recurrent state is 118 saved for the next time step. 119 120 Parameters 121 ---------- 122 *args : tuple 123 Positional arguments containing the dendrite output. 124 **kwargs : dict 125 Keyword arguments from the dendrite output. 126 127 Returns 128 ------- 129 Any 130 The filtered dendrite output to be added to the neuron output. 131 """ 132 pass 133 134 @abstractmethod 135 def clear_processor(self): 136 """ 137 Clear all internal processor state. 138 139 Resets the processor by removing all stored state variables. Must 140 be called before saving or safe_tensors will run into errors. 141 Implementations should safely check for attribute existence before 142 deletion to avoid errors. 143 """ 144 pass 145 146 147# General multi output processor for any number that ignores later ones 148class MultiOutputProcessor: 149 """Processor for handling multiple outputs, ignoring later ones.""" 150 151 def post_n1(self, *args, **kwargs): 152 """Saves extra outputs and returns the first output. 153 154 Parameters 155 ---------- 156 *args : tuple 157 Contains the modules output tuple. 158 **kwargs : dict 159 Unused keyword arguments. 160 161 Returns 162 ------- 163 torch.Tensor 164 The first tensor of the tuple 165 """ 166 out = args[0][0] 167 extra_out = args[0][1:] 168 self.extra_out = extra_out 169 return out 170 171 def post_n2(self, *args, **kwargs): 172 """Combine output with stored extra outputs. 173 174 Parameters 175 ---------- 176 *args : torch.tensor 177 The first tensor combined with dendrite output. 178 **kwargs : dict 179 Unused keyword arguments. 180 181 Returns 182 ------- 183 tuple 184 The recombined output tuple wth the new first output modified 185 """ 186 out = args[0] 187 if isinstance(self.extra_out, tuple): 188 return (out,) + self.extra_out 189 else: 190 return (out,) + (self.extra_out,) 191 192 def pre_d(self, *args, **kwargs): 193 """Pass through arguments unchanged for dendrite preprocessing. 194 195 Parameters 196 ---------- 197 *args : tuple 198 Positional arguments containing inputs to the PAI module. 199 **kwargs : dict 200 Keyword arguments containing inputs to the PAI module. 201 202 Returns 203 ------- 204 args : tuple 205 Positional arguments containing inputs to the PAI module. 206 kwargs : dict 207 Keyword arguments containing inputs to the PAI module. 208 """ 209 return args, kwargs 210 211 def post_d(self, *args, **kwargs): 212 """Extract first output for dendrite postprocessing. 213 214 Parameters 215 ---------- 216 *args : tuple 217 Contains the dendrite modules output tuple. 218 **kwargs : dict 219 Unused keyword arguments. 220 221 Returns 222 ------- 223 torch.Tensor 224 The first tensor of the tuple 225 """ 226 out = args[0][0] 227 return out 228 229 def clear_processor(self): 230 """Clear stored processor state.""" 231 232 if hasattr(self, "extra_out"): 233 delattr(self, "extra_out") 234 235class LSTMCellProcessor(PAIProcessor): 236 """Processor for LSTM cells to handle hidden and cell states.""" 237 238 def post_n1(self, *args, **kwargs): 239 """ 240 Extract hidden state from LSTM output for dendrite processing. 241 242 Separates the hidden state (h_t) from the cell state (c_t) in the 243 LSTM output tuple. Stores the cell state temporarily since only the 244 hidden state should be modified by dendrites. 245 246 Parameters 247 ---------- 248 *args : tuple 249 Contains LSTM output tuple (h_t, c_t) as first element. 250 **kwargs : dict 251 Unused keyword arguments. 252 253 Returns 254 ------- 255 torch.Tensor 256 Hidden state h_t to be passed to dendrite processing. 257 """ 258 h_t = args[0][0] 259 c_t = args[0][1] 260 # Store the cell state temporarily and just use the hidden state 261 # to do Dendrite functions 262 self.c_t_n = c_t 263 return h_t 264 265 def post_n2(self, *args, **kwargs): 266 """ 267 Recombine dendrite-modified hidden state with cell state. 268 269 Takes the hidden state that has been modified by dendrite operations 270 and combines it with the stored cell state to produce the complete 271 LSTM output tuple. 272 273 Parameters 274 ---------- 275 *args : tuple 276 Contains the dendrite-modified hidden state h_t. 277 **kwargs : dict 278 Unused keyword arguments. 279 280 Returns 281 ------- 282 tuple 283 Complete LSTM output (h_t, c_t) where h_t has been modified. 284 """ 285 h_t = args[0] 286 return h_t, self.c_t_n 287 288 def pre_d(self, *args, **kwargs): 289 """ 290 Filter LSTMCell input for dendrite based on initialization state. 291 292 Checks if this is the first time step (all zeros in h_t) or a 293 subsequent step. For the first step, passes through the original 294 inputs. For subsequent steps, replaces the neuron's hidden state 295 with the dendrite's own internal state from the previous iteration. 296 297 Parameters 298 ---------- 299 *args : tuple 300 Contains (input, (h_t, c_t)) where input is the external input 301 and (h_t, c_t) is the neuron's recurrent state. 302 **kwargs : dict 303 Keyword arguments to pass through. 304 305 Returns 306 ------- 307 tuple 308 ((processed_input, processed_state), kwargs) for dendrite call. 309 """ 310 h_t = args[1][0] 311 # If its the initial step then just use the normal input and zeros 312 if h_t.sum() == 0: 313 return args, kwargs 314 # If its not the first one then return the input it got with its own 315 # h_t and c_t to replace neurons 316 else: 317 return (args[0], (self.h_t_d, self.c_t_d)), kwargs 318 319 def post_d(self, *args, **kwargs): 320 """ 321 Extract and store dendrite's LSTM state for next iteration. 322 323 Separates the dendrite's hidden and cell states from its output tuple, 324 stores both for use in the next time step, and returns only the hidden 325 state to be combined with the neuron's output. 326 327 Parameters 328 ---------- 329 *args : tuple 330 Contains dendrite LSTM output tuple (h_t, c_t). 331 **kwargs : dict 332 Unused keyword arguments. 333 334 Returns 335 ------- 336 torch.Tensor 337 Hidden state h_t to be added to the neuron output. 338 """ 339 h_t = args[0][0] 340 c_t = args[0][1] 341 self.h_t_d = h_t 342 self.c_t_d = c_t 343 return h_t 344 345 def clear_processor(self): 346 """ 347 Clear all stored LSTM states. 348 349 Removes dendrite hidden state (h_t_d), dendrite cell state (c_t_d), 350 and temporarily stored neuron cell state (c_t_n). Safe to call even 351 if attributes don't exist. 352 """ 353 if hasattr(self, "h_t_d"): 354 delattr(self, "h_t_d") 355 if hasattr(self, "c_t_d"): 356 delattr(self, "c_t_d") 357 if hasattr(self, "c_t_n"): 358 delattr(self, "c_t_n") 359 360 361 362class LSTMProcessor(PAIProcessor): 363 """Processor for LSTM to handle hidden and output states.""" 364 365 def post_n1(self, *args, **kwargs): 366 """ 367 Extract hidden state from LSTM output for dendrite processing. 368 369 Separates the hidden state from the output in the 370 LSTM output tuple. Stores the hidden state temporarily since only the 371 output state should be modified by dendrites. 372 373 Parameters 374 ---------- 375 *args : tuple 376 Contains LSTM output tuple (output, hidden) as first element. 377 **kwargs : dict 378 Unused keyword arguments. 379 380 Returns 381 ------- 382 torch.Tensor 383 Output state to be passed to dendrite processing. 384 """ 385 output = args[0][0] 386 hidden = args[0][1] 387 # Store the hidden state temporarily and just use the output state 388 # to do Dendrite functions 389 self.hidden_n = hidden 390 return output 391 392 def post_n2(self, *args, **kwargs): 393 """ 394 Recombine dendrite-modified output with hidden tuple. 395 396 Takes the output state that has been modified by dendrite operations 397 and combines it with the stored hidden state to produce the complete 398 LSTM output tuple. 399 400 Parameters 401 ---------- 402 *args : tuple 403 Contains the dendrite-modified output state. 404 **kwargs : dict 405 Unused keyword arguments. 406 407 Returns 408 ------- 409 tuple 410 Complete LSTM output (output, hidden) where output has been modified. 411 """ 412 output = args[0] 413 return output, self.hidden_n 414 415 def pre_d(self, *args, **kwargs): 416 """ 417 LSTM input is just the tensor which also goes to the dendrite 418 419 Parameters 420 ---------- 421 *args : 422 Input tensor 423 **kwargs : dict 424 Empty 425 426 Returns 427 ------- 428 tuple 429 (output, hidden) 430 """ 431 return args, kwargs 432 433 def post_d(self, *args, **kwargs): 434 """ 435 Extract dendrite's output to combine. 436 437 Parameters 438 ---------- 439 *args : tuple 440 Contains dendrite LSTM output tuple (output, hidden). 441 **kwargs : dict 442 Unused keyword arguments. 443 444 Returns 445 ------- 446 torch.Tensor 447 Output state to be added to the neuron output. 448 """ 449 output = args[0][0] 450 hidden = args[0][1] 451 return output 452 453 def clear_processor(self): 454 """ 455 Clear all stored LSTM states. 456 457 """ 458 if hasattr(self, "hidden_n"): 459 delattr(self, "hidden_n") 460 461 462class LSTMProcessorLastHidden(PAIProcessor): 463 """Processor for LSTM to forward the last hidden.""" 464 465 def post_n1(self, *args, **kwargs): 466 """ 467 Extract the last hidden to combine with dendrites 468 469 Parameters 470 ---------- 471 *args : tuple 472 Contains LSTM output tuple (output, hidden) as first element. 473 **kwargs : dict 474 Unused keyword arguments. 475 476 Returns 477 ------- 478 torch.Tensor 479 Output state to be passed to dendrite processing. 480 """ 481 ignored_output = args[0][0] 482 last_hidden = args[0][1][-1] 483 484 return last_hidden 485 486 def post_n2(self, *args, **kwargs): 487 """ 488 Recombine dendrite-modified last hidden, and append None just to maintain output format 489 490 Parameters 491 ---------- 492 *args : tuple 493 Contains the dendrite-modified output state. 494 **kwargs : dict 495 Unused keyword arguments. 496 497 Returns 498 ------- 499 tuple 500 Complete LSTM output (output, hidden) where output has been modified. 501 """ 502 combined_last_hidden = args[0] 503 return None, combined_last_hidden 504 505 def pre_d(self, *args, **kwargs): 506 """ 507 LSTM input is just the tensor which also goes to the dendrite 508 509 Parameters 510 ---------- 511 *args : 512 Input tensor 513 **kwargs : dict 514 Empty 515 516 Returns 517 ------- 518 tuple 519 (output, hidden) 520 """ 521 return args, kwargs 522 523 def post_d(self, *args, **kwargs): 524 """ 525 Extract extract the dendrites last hidden to combine with neurons. 526 527 Parameters 528 ---------- 529 *args : tuple 530 Contains dendrite LSTM output tuple (output, hidden). 531 **kwargs : dict 532 Unused keyword arguments. 533 534 Returns 535 ------- 536 torch.Tensor 537 Output state to be added to the neuron output. 538 """ 539 ignored_output = args[0][0] 540 last_hidden = args[0][1][-1] 541 return last_hidden 542 543 def clear_processor(self): 544 # Nothing is stored 545 pass 546 547class ResNetPAI(nn.Module): 548 """PB-compatible ResNet wrapper. 549 550 All normalization layers should be wrapped in a PAISequential, or other 551 wrapped module. When working with a predefined model the following shows 552 an example of how to create a module for modules_to_replace. 553 """ 554 555 def __init__(self, other_resnet): 556 """Initialize ResNetPAI from existing ResNet model. 557 558 Parameters 559 ---------- 560 *args : other_resnet : torchvision.models.resnet.ResNet 561 An existing ResNet model to convert to PAI-compatible format. 562 """ 563 super(ResNetPAI, self).__init__() 564 565 # For the most part, just copy the exact values from the original module 566 self._norm_layer = other_resnet._norm_layer 567 self.inplanes = other_resnet.inplanes 568 self.dilation = other_resnet.dilation 569 self.groups = other_resnet.groups 570 self.base_width = other_resnet.base_width 571 572 # For the component to be changed, define a PAISequential with the old 573 # modules included 574 self.b1 = GPA.PAISequential([other_resnet.conv1, other_resnet.bn1]) 575 576 self.relu = other_resnet.relu 577 self.maxpool = other_resnet.maxpool 578 579 for i in range(1, 5): 580 layer_name = "layer" + str(i) 581 original_layer = getattr(other_resnet, layer_name) 582 pb_layer = self._make_layer_pb(original_layer, other_resnet, i) 583 setattr(self, layer_name, pb_layer) 584 585 self.avgpool = other_resnet.avgpool 586 self.fc = other_resnet.fc 587 588 def _make_layer_pb(self, other_block_set, other_resnet, block_id): 589 """Convert ResNet layer blocks to PB-compatible format. 590 591 Parameters 592 ---------- 593 other_block_set : torch.vision.models.resnet.any_block 594 A set of blocks from the original ResNet model. 595 other_resnet : torchvision.models.resnet.ResNet 596 The original ResNet model. 597 block_id : int 598 The layer number being converted. 599 Returns 600 ------- 601 nn.Sequential 602 A sequential container with the converted blocks. 603 """ 604 layers = [] 605 for i in range(len(other_block_set)): 606 block_type = type(other_block_set[i]) 607 if block_type == resnet_pt.BasicBlock: 608 layers.append(other_block_set[i]) 609 elif block_type == resnet_pt.Bottleneck: 610 layers.append(other_block_set[i]) 611 else: 612 print( 613 "Your resnet uses a block type that has not been " 614 "accounted for. Customization might be required." 615 ) 616 layer_name = "layer" + str(block_id) 617 print(type(getattr(other_resnet, layer_name))) 618 pdb.set_trace() 619 return nn.Sequential(*layers) 620 621 def _forward_impl(self, x): 622 """Implementation of the forward pass. 623 624 Parameters 625 ---------- 626 x : torch.Tensor 627 Input tensor to the network. 628 629 Returns 630 ------- 631 torch.Tensor 632 Output tensor from the network. 633 """ 634 # Modified b1 rather than conv1 and bn1 635 x = self.b1(x) 636 # Rest of forward remains the same 637 x = F.relu(x) 638 x = self.maxpool(x) 639 640 x = self.layer1(x) 641 x = self.layer2(x) 642 x = self.layer3(x) 643 x = self.layer4(x) 644 645 x = self.avgpool(x) 646 x = torch.flatten(x, 1) 647 x = self.fc(x) 648 649 return x 650 651 def forward(self, x): 652 """Forward pass through the network. 653 654 Parameters 655 ---------- 656 x : torch.Tensor 657 Input tensor to the network. 658 659 Returns 660 ------- 661 torch.Tensor 662 Output tensor from the network. 663 """ 664 return self._forward_impl(x) 665 666 667class ResNetPAIPreFC(nn.Module): 668 """PB-compatible ResNet wrapper. 669 670 All normalization layers should be wrapped in a PAISequential, or other 671 wrapped module. When working with a predefined model the following shows 672 an example of how to create a module for modules_to_replace. 673 """ 674 675 def __init__(self, other_resnet): 676 """Initialize ResNetPAI from existing ResNet model. 677 678 Parameters 679 ---------- 680 *args : other_resnet : torchvision.models.resnet.ResNet 681 An existing ResNet model to convert to PAI-compatible format. 682 """ 683 super(ResNetPAIPreFC, self).__init__() 684 685 # For the most part, just copy the exact values from the original module 686 self._norm_layer = other_resnet._norm_layer 687 self.inplanes = other_resnet.inplanes 688 self.dilation = other_resnet.dilation 689 self.groups = other_resnet.groups 690 self.base_width = other_resnet.base_width 691 692 # For the component to be changed, define a PAISequential with the old 693 # modules included 694 self.conv1 = other_resnet.conv1 695 self.bn1 = other_resnet.bn1 696 697 self.relu = other_resnet.relu 698 self.maxpool = other_resnet.maxpool 699 700 for i in range(1, 5): 701 layer_name = "layer" + str(i) 702 original_layer = getattr(other_resnet, layer_name) 703 pb_layer = self._make_layer_pb(original_layer, other_resnet, i) 704 setattr(self, layer_name, pb_layer) 705 706 self.avgpool = other_resnet.avgpool 707 708 # Create pre_fc layer with dimensions matching layer4 output (same as fc input) 709 fc_in_features = other_resnet.fc.in_features 710 self.pre_fc = nn.Linear(fc_in_features, fc_in_features) 711 712 self.fc = other_resnet.fc 713 714 def _make_layer_pb(self, other_block_set, other_resnet, block_id): 715 """Convert ResNet layer blocks to PB-compatible format. 716 717 Parameters 718 ---------- 719 other_block_set : torch.vision.models.resnet.any_block 720 A set of blocks from the original ResNet model. 721 other_resnet : torchvision.models.resnet.ResNet 722 The original ResNet model. 723 block_id : int 724 The layer number being converted. 725 Returns 726 ------- 727 nn.Sequential 728 A sequential container with the converted blocks. 729 """ 730 layers = [] 731 for i in range(len(other_block_set)): 732 block_type = type(other_block_set[i]) 733 if block_type == resnet_pt.BasicBlock: 734 layers.append(other_block_set[i]) 735 elif block_type == resnet_pt.Bottleneck: 736 layers.append(other_block_set[i]) 737 else: 738 print( 739 "Your resnet uses a block type that has not been " 740 "accounted for. Customization might be required." 741 ) 742 layer_name = "layer" + str(block_id) 743 print(type(getattr(other_resnet, layer_name))) 744 pdb.set_trace() 745 return nn.Sequential(*layers) 746 747 def _forward_impl(self, x): 748 """Implementation of the forward pass. 749 750 Parameters 751 ---------- 752 x : torch.Tensor 753 Input tensor to the network. 754 755 Returns 756 ------- 757 torch.Tensor 758 Output tensor from the network. 759 """ 760 # Modified b1 rather than conv1 and bn1 761 x = self.conv1(x) 762 x = self.bn1(x) 763 # Rest of forward remains the same 764 x = F.relu(x) 765 x = self.maxpool(x) 766 767 x = self.layer1(x) 768 x = self.layer2(x) 769 x = self.layer3(x) 770 x = self.layer4(x) 771 772 x = self.avgpool(x) 773 x = torch.flatten(x, 1) 774 x = self.pre_fc(x) 775 x = F.relu(x) 776 x = self.fc(x) 777 778 return x 779 780 def forward(self, x): 781 """Forward pass through the network. 782 783 Parameters 784 ---------- 785 x : torch.Tensor 786 Input tensor to the network. 787 788 Returns 789 ------- 790 torch.Tensor 791 Output tensor from the network. 792 """ 793 return self._forward_impl(x)
class
PAIProcessor(abc.ABC):
28class PAIProcessor(ABC): 29 """ 30 Abstract base class for processing neuron and dendrite operations. 31 32 Processors handle state management and data flow between neurons and 33 dendrites, allowing for custom pre/post processing of modules which have 34 multiple inputs and outputs, rather than the default single tensor input/output. 35 Subclasses should implement the five core processing methods to handle 36 their specific state management needs. 37 """ 38 39 @abstractmethod 40 def post_n1(self, *args, **kwargs): 41 """ 42 Post-process neuron output before dendrite processing. 43 44 Called immediately after the main module/neuron is executed and before 45 any dendrite processing occurs. This method should extract and return 46 only the tensor of the neuron output that should be seen by 47 dendrite operations. 48 49 Parameters 50 ---------- 51 *args : tuple 52 Positional arguments, typically containing the neuron output. 53 **kwargs : dict 54 Keyword arguments from the neuron output. 55 56 Returns 57 ------- 58 Any 59 The filtered output to be passed to dendrite processing. 60 """ 61 pass 62 63 @abstractmethod 64 def post_n2(self, *args, **kwargs): 65 """ 66 Post-process dendrite-modified output before final return. 67 68 Called after dendrite processing is complete and before passing the 69 final value forward in the network. This method should combine the 70 dendrite-modified output with any stored state to produce the complete 71 output that matches the expected format of the main module. 72 73 Parameters 74 ---------- 75 *args : tuple 76 Positional arguments containing the dendrite-modified output. 77 **kwargs : dict 78 Keyword arguments from the processing chain. 79 80 Returns 81 ------- 82 Any 83 The complete output in the format expected by downstream components. 84 """ 85 pass 86 87 @abstractmethod 88 def pre_d(self, *args, **kwargs): 89 """ 90 Pre-process input before dendrite operations. 91 92 Filters and prepares inputs for dendrite processing. This method handles 93 special cases such as initial time steps vs. subsequent iterations, 94 ensuring dendrites receive the appropriate inputs (e.g., external inputs 95 vs. internal recurrent state). 96 97 Parameters 98 ---------- 99 *args : tuple 100 Positional arguments containing inputs to the PAI module. 101 **kwargs : dict 102 Keyword arguments containing inputs to the PAI module. 103 104 Returns 105 ------- 106 tuple 107 A tuple of (processed_args, processed_kwargs) to pass to dendrite. 108 """ 109 pass 110 111 @abstractmethod 112 def post_d(self, *args, **kwargs): 113 """ 114 Post-process dendrite output and manage state. 115 116 Processes the output from dendrite operations, storing any state needed 117 for future iterations and returning only the portion that should be 118 combined with the neuron output. E.g. this is where recurrent state is 119 saved for the next time step. 120 121 Parameters 122 ---------- 123 *args : tuple 124 Positional arguments containing the dendrite output. 125 **kwargs : dict 126 Keyword arguments from the dendrite output. 127 128 Returns 129 ------- 130 Any 131 The filtered dendrite output to be added to the neuron output. 132 """ 133 pass 134 135 @abstractmethod 136 def clear_processor(self): 137 """ 138 Clear all internal processor state. 139 140 Resets the processor by removing all stored state variables. Must 141 be called before saving or safe_tensors will run into errors. 142 Implementations should safely check for attribute existence before 143 deletion to avoid errors. 144 """ 145 pass
Abstract base class for processing neuron and dendrite operations.
Processors handle state management and data flow between neurons and dendrites, allowing for custom pre/post processing of modules which have multiple inputs and outputs, rather than the default single tensor input/output. Subclasses should implement the five core processing methods to handle their specific state management needs.
@abstractmethod
def
post_n1(self, *args, **kwargs):
39 @abstractmethod 40 def post_n1(self, *args, **kwargs): 41 """ 42 Post-process neuron output before dendrite processing. 43 44 Called immediately after the main module/neuron is executed and before 45 any dendrite processing occurs. This method should extract and return 46 only the tensor of the neuron output that should be seen by 47 dendrite operations. 48 49 Parameters 50 ---------- 51 *args : tuple 52 Positional arguments, typically containing the neuron output. 53 **kwargs : dict 54 Keyword arguments from the neuron output. 55 56 Returns 57 ------- 58 Any 59 The filtered output to be passed to dendrite processing. 60 """ 61 pass
Post-process neuron output before dendrite processing.
Called immediately after the main module/neuron is executed and before any dendrite processing occurs. This method should extract and return only the tensor of the neuron output that should be seen by dendrite operations.
Parameters
- *args (tuple): Positional arguments, typically containing the neuron output.
- **kwargs (dict): Keyword arguments from the neuron output.
Returns
- Any: The filtered output to be passed to dendrite processing.
@abstractmethod
def
post_n2(self, *args, **kwargs):
63 @abstractmethod 64 def post_n2(self, *args, **kwargs): 65 """ 66 Post-process dendrite-modified output before final return. 67 68 Called after dendrite processing is complete and before passing the 69 final value forward in the network. This method should combine the 70 dendrite-modified output with any stored state to produce the complete 71 output that matches the expected format of the main module. 72 73 Parameters 74 ---------- 75 *args : tuple 76 Positional arguments containing the dendrite-modified output. 77 **kwargs : dict 78 Keyword arguments from the processing chain. 79 80 Returns 81 ------- 82 Any 83 The complete output in the format expected by downstream components. 84 """ 85 pass
Post-process dendrite-modified output before final return.
Called after dendrite processing is complete and before passing the final value forward in the network. This method should combine the dendrite-modified output with any stored state to produce the complete output that matches the expected format of the main module.
Parameters
- *args (tuple): Positional arguments containing the dendrite-modified output.
- **kwargs (dict): Keyword arguments from the processing chain.
Returns
- Any: The complete output in the format expected by downstream components.
@abstractmethod
def
pre_d(self, *args, **kwargs):
87 @abstractmethod 88 def pre_d(self, *args, **kwargs): 89 """ 90 Pre-process input before dendrite operations. 91 92 Filters and prepares inputs for dendrite processing. This method handles 93 special cases such as initial time steps vs. subsequent iterations, 94 ensuring dendrites receive the appropriate inputs (e.g., external inputs 95 vs. internal recurrent state). 96 97 Parameters 98 ---------- 99 *args : tuple 100 Positional arguments containing inputs to the PAI module. 101 **kwargs : dict 102 Keyword arguments containing inputs to the PAI module. 103 104 Returns 105 ------- 106 tuple 107 A tuple of (processed_args, processed_kwargs) to pass to dendrite. 108 """ 109 pass
Pre-process input before dendrite operations.
Filters and prepares inputs for dendrite processing. This method handles special cases such as initial time steps vs. subsequent iterations, ensuring dendrites receive the appropriate inputs (e.g., external inputs vs. internal recurrent state).
Parameters
- *args (tuple): Positional arguments containing inputs to the PAI module.
- **kwargs (dict): Keyword arguments containing inputs to the PAI module.
Returns
- tuple: A tuple of (processed_args, processed_kwargs) to pass to dendrite.
@abstractmethod
def
post_d(self, *args, **kwargs):
111 @abstractmethod 112 def post_d(self, *args, **kwargs): 113 """ 114 Post-process dendrite output and manage state. 115 116 Processes the output from dendrite operations, storing any state needed 117 for future iterations and returning only the portion that should be 118 combined with the neuron output. E.g. this is where recurrent state is 119 saved for the next time step. 120 121 Parameters 122 ---------- 123 *args : tuple 124 Positional arguments containing the dendrite output. 125 **kwargs : dict 126 Keyword arguments from the dendrite output. 127 128 Returns 129 ------- 130 Any 131 The filtered dendrite output to be added to the neuron output. 132 """ 133 pass
Post-process dendrite output and manage state.
Processes the output from dendrite operations, storing any state needed for future iterations and returning only the portion that should be combined with the neuron output. E.g. this is where recurrent state is saved for the next time step.
Parameters
- *args (tuple): Positional arguments containing the dendrite output.
- **kwargs (dict): Keyword arguments from the dendrite output.
Returns
- Any: The filtered dendrite output to be added to the neuron output.
@abstractmethod
def
clear_processor(self):
135 @abstractmethod 136 def clear_processor(self): 137 """ 138 Clear all internal processor state. 139 140 Resets the processor by removing all stored state variables. Must 141 be called before saving or safe_tensors will run into errors. 142 Implementations should safely check for attribute existence before 143 deletion to avoid errors. 144 """ 145 pass
Clear all internal processor state.
Resets the processor by removing all stored state variables. Must be called before saving or safe_tensors will run into errors. Implementations should safely check for attribute existence before deletion to avoid errors.
class
MultiOutputProcessor:
149class MultiOutputProcessor: 150 """Processor for handling multiple outputs, ignoring later ones.""" 151 152 def post_n1(self, *args, **kwargs): 153 """Saves extra outputs and returns the first output. 154 155 Parameters 156 ---------- 157 *args : tuple 158 Contains the modules output tuple. 159 **kwargs : dict 160 Unused keyword arguments. 161 162 Returns 163 ------- 164 torch.Tensor 165 The first tensor of the tuple 166 """ 167 out = args[0][0] 168 extra_out = args[0][1:] 169 self.extra_out = extra_out 170 return out 171 172 def post_n2(self, *args, **kwargs): 173 """Combine output with stored extra outputs. 174 175 Parameters 176 ---------- 177 *args : torch.tensor 178 The first tensor combined with dendrite output. 179 **kwargs : dict 180 Unused keyword arguments. 181 182 Returns 183 ------- 184 tuple 185 The recombined output tuple wth the new first output modified 186 """ 187 out = args[0] 188 if isinstance(self.extra_out, tuple): 189 return (out,) + self.extra_out 190 else: 191 return (out,) + (self.extra_out,) 192 193 def pre_d(self, *args, **kwargs): 194 """Pass through arguments unchanged for dendrite preprocessing. 195 196 Parameters 197 ---------- 198 *args : tuple 199 Positional arguments containing inputs to the PAI module. 200 **kwargs : dict 201 Keyword arguments containing inputs to the PAI module. 202 203 Returns 204 ------- 205 args : tuple 206 Positional arguments containing inputs to the PAI module. 207 kwargs : dict 208 Keyword arguments containing inputs to the PAI module. 209 """ 210 return args, kwargs 211 212 def post_d(self, *args, **kwargs): 213 """Extract first output for dendrite postprocessing. 214 215 Parameters 216 ---------- 217 *args : tuple 218 Contains the dendrite modules output tuple. 219 **kwargs : dict 220 Unused keyword arguments. 221 222 Returns 223 ------- 224 torch.Tensor 225 The first tensor of the tuple 226 """ 227 out = args[0][0] 228 return out 229 230 def clear_processor(self): 231 """Clear stored processor state.""" 232 233 if hasattr(self, "extra_out"): 234 delattr(self, "extra_out")
Processor for handling multiple outputs, ignoring later ones.
def
post_n1(self, *args, **kwargs):
152 def post_n1(self, *args, **kwargs): 153 """Saves extra outputs and returns the first output. 154 155 Parameters 156 ---------- 157 *args : tuple 158 Contains the modules output tuple. 159 **kwargs : dict 160 Unused keyword arguments. 161 162 Returns 163 ------- 164 torch.Tensor 165 The first tensor of the tuple 166 """ 167 out = args[0][0] 168 extra_out = args[0][1:] 169 self.extra_out = extra_out 170 return out
Saves extra outputs and returns the first output.
Parameters
- *args (tuple): Contains the modules output tuple.
- **kwargs (dict): Unused keyword arguments.
Returns
- torch.Tensor: The first tensor of the tuple
def
post_n2(self, *args, **kwargs):
172 def post_n2(self, *args, **kwargs): 173 """Combine output with stored extra outputs. 174 175 Parameters 176 ---------- 177 *args : torch.tensor 178 The first tensor combined with dendrite output. 179 **kwargs : dict 180 Unused keyword arguments. 181 182 Returns 183 ------- 184 tuple 185 The recombined output tuple wth the new first output modified 186 """ 187 out = args[0] 188 if isinstance(self.extra_out, tuple): 189 return (out,) + self.extra_out 190 else: 191 return (out,) + (self.extra_out,)
Combine output with stored extra outputs.
Parameters
- *args (torch.tensor): The first tensor combined with dendrite output.
- **kwargs (dict): Unused keyword arguments.
Returns
- tuple: The recombined output tuple wth the new first output modified
def
pre_d(self, *args, **kwargs):
193 def pre_d(self, *args, **kwargs): 194 """Pass through arguments unchanged for dendrite preprocessing. 195 196 Parameters 197 ---------- 198 *args : tuple 199 Positional arguments containing inputs to the PAI module. 200 **kwargs : dict 201 Keyword arguments containing inputs to the PAI module. 202 203 Returns 204 ------- 205 args : tuple 206 Positional arguments containing inputs to the PAI module. 207 kwargs : dict 208 Keyword arguments containing inputs to the PAI module. 209 """ 210 return args, kwargs
Pass through arguments unchanged for dendrite preprocessing.
Parameters
- *args (tuple): Positional arguments containing inputs to the PAI module.
- **kwargs (dict): Keyword arguments containing inputs to the PAI module.
Returns
- args (tuple): Positional arguments containing inputs to the PAI module.
- kwargs (dict): Keyword arguments containing inputs to the PAI module.
def
post_d(self, *args, **kwargs):
212 def post_d(self, *args, **kwargs): 213 """Extract first output for dendrite postprocessing. 214 215 Parameters 216 ---------- 217 *args : tuple 218 Contains the dendrite modules output tuple. 219 **kwargs : dict 220 Unused keyword arguments. 221 222 Returns 223 ------- 224 torch.Tensor 225 The first tensor of the tuple 226 """ 227 out = args[0][0] 228 return out
Extract first output for dendrite postprocessing.
Parameters
- *args (tuple): Contains the dendrite modules output tuple.
- **kwargs (dict): Unused keyword arguments.
Returns
- torch.Tensor: The first tensor of the tuple
236class LSTMCellProcessor(PAIProcessor): 237 """Processor for LSTM cells to handle hidden and cell states.""" 238 239 def post_n1(self, *args, **kwargs): 240 """ 241 Extract hidden state from LSTM output for dendrite processing. 242 243 Separates the hidden state (h_t) from the cell state (c_t) in the 244 LSTM output tuple. Stores the cell state temporarily since only the 245 hidden state should be modified by dendrites. 246 247 Parameters 248 ---------- 249 *args : tuple 250 Contains LSTM output tuple (h_t, c_t) as first element. 251 **kwargs : dict 252 Unused keyword arguments. 253 254 Returns 255 ------- 256 torch.Tensor 257 Hidden state h_t to be passed to dendrite processing. 258 """ 259 h_t = args[0][0] 260 c_t = args[0][1] 261 # Store the cell state temporarily and just use the hidden state 262 # to do Dendrite functions 263 self.c_t_n = c_t 264 return h_t 265 266 def post_n2(self, *args, **kwargs): 267 """ 268 Recombine dendrite-modified hidden state with cell state. 269 270 Takes the hidden state that has been modified by dendrite operations 271 and combines it with the stored cell state to produce the complete 272 LSTM output tuple. 273 274 Parameters 275 ---------- 276 *args : tuple 277 Contains the dendrite-modified hidden state h_t. 278 **kwargs : dict 279 Unused keyword arguments. 280 281 Returns 282 ------- 283 tuple 284 Complete LSTM output (h_t, c_t) where h_t has been modified. 285 """ 286 h_t = args[0] 287 return h_t, self.c_t_n 288 289 def pre_d(self, *args, **kwargs): 290 """ 291 Filter LSTMCell input for dendrite based on initialization state. 292 293 Checks if this is the first time step (all zeros in h_t) or a 294 subsequent step. For the first step, passes through the original 295 inputs. For subsequent steps, replaces the neuron's hidden state 296 with the dendrite's own internal state from the previous iteration. 297 298 Parameters 299 ---------- 300 *args : tuple 301 Contains (input, (h_t, c_t)) where input is the external input 302 and (h_t, c_t) is the neuron's recurrent state. 303 **kwargs : dict 304 Keyword arguments to pass through. 305 306 Returns 307 ------- 308 tuple 309 ((processed_input, processed_state), kwargs) for dendrite call. 310 """ 311 h_t = args[1][0] 312 # If its the initial step then just use the normal input and zeros 313 if h_t.sum() == 0: 314 return args, kwargs 315 # If its not the first one then return the input it got with its own 316 # h_t and c_t to replace neurons 317 else: 318 return (args[0], (self.h_t_d, self.c_t_d)), kwargs 319 320 def post_d(self, *args, **kwargs): 321 """ 322 Extract and store dendrite's LSTM state for next iteration. 323 324 Separates the dendrite's hidden and cell states from its output tuple, 325 stores both for use in the next time step, and returns only the hidden 326 state to be combined with the neuron's output. 327 328 Parameters 329 ---------- 330 *args : tuple 331 Contains dendrite LSTM output tuple (h_t, c_t). 332 **kwargs : dict 333 Unused keyword arguments. 334 335 Returns 336 ------- 337 torch.Tensor 338 Hidden state h_t to be added to the neuron output. 339 """ 340 h_t = args[0][0] 341 c_t = args[0][1] 342 self.h_t_d = h_t 343 self.c_t_d = c_t 344 return h_t 345 346 def clear_processor(self): 347 """ 348 Clear all stored LSTM states. 349 350 Removes dendrite hidden state (h_t_d), dendrite cell state (c_t_d), 351 and temporarily stored neuron cell state (c_t_n). Safe to call even 352 if attributes don't exist. 353 """ 354 if hasattr(self, "h_t_d"): 355 delattr(self, "h_t_d") 356 if hasattr(self, "c_t_d"): 357 delattr(self, "c_t_d") 358 if hasattr(self, "c_t_n"): 359 delattr(self, "c_t_n")
Processor for LSTM cells to handle hidden and cell states.
def
post_n1(self, *args, **kwargs):
239 def post_n1(self, *args, **kwargs): 240 """ 241 Extract hidden state from LSTM output for dendrite processing. 242 243 Separates the hidden state (h_t) from the cell state (c_t) in the 244 LSTM output tuple. Stores the cell state temporarily since only the 245 hidden state should be modified by dendrites. 246 247 Parameters 248 ---------- 249 *args : tuple 250 Contains LSTM output tuple (h_t, c_t) as first element. 251 **kwargs : dict 252 Unused keyword arguments. 253 254 Returns 255 ------- 256 torch.Tensor 257 Hidden state h_t to be passed to dendrite processing. 258 """ 259 h_t = args[0][0] 260 c_t = args[0][1] 261 # Store the cell state temporarily and just use the hidden state 262 # to do Dendrite functions 263 self.c_t_n = c_t 264 return h_t
Extract hidden state from LSTM output for dendrite processing.
Separates the hidden state (h_t) from the cell state (c_t) in the LSTM output tuple. Stores the cell state temporarily since only the hidden state should be modified by dendrites.
Parameters
- *args (tuple): Contains LSTM output tuple (h_t, c_t) as first element.
- **kwargs (dict): Unused keyword arguments.
Returns
- torch.Tensor: Hidden state h_t to be passed to dendrite processing.
def
post_n2(self, *args, **kwargs):
266 def post_n2(self, *args, **kwargs): 267 """ 268 Recombine dendrite-modified hidden state with cell state. 269 270 Takes the hidden state that has been modified by dendrite operations 271 and combines it with the stored cell state to produce the complete 272 LSTM output tuple. 273 274 Parameters 275 ---------- 276 *args : tuple 277 Contains the dendrite-modified hidden state h_t. 278 **kwargs : dict 279 Unused keyword arguments. 280 281 Returns 282 ------- 283 tuple 284 Complete LSTM output (h_t, c_t) where h_t has been modified. 285 """ 286 h_t = args[0] 287 return h_t, self.c_t_n
Recombine dendrite-modified hidden state with cell state.
Takes the hidden state that has been modified by dendrite operations and combines it with the stored cell state to produce the complete LSTM output tuple.
Parameters
- *args (tuple): Contains the dendrite-modified hidden state h_t.
- **kwargs (dict): Unused keyword arguments.
Returns
- tuple: Complete LSTM output (h_t, c_t) where h_t has been modified.
def
pre_d(self, *args, **kwargs):
289 def pre_d(self, *args, **kwargs): 290 """ 291 Filter LSTMCell input for dendrite based on initialization state. 292 293 Checks if this is the first time step (all zeros in h_t) or a 294 subsequent step. For the first step, passes through the original 295 inputs. For subsequent steps, replaces the neuron's hidden state 296 with the dendrite's own internal state from the previous iteration. 297 298 Parameters 299 ---------- 300 *args : tuple 301 Contains (input, (h_t, c_t)) where input is the external input 302 and (h_t, c_t) is the neuron's recurrent state. 303 **kwargs : dict 304 Keyword arguments to pass through. 305 306 Returns 307 ------- 308 tuple 309 ((processed_input, processed_state), kwargs) for dendrite call. 310 """ 311 h_t = args[1][0] 312 # If its the initial step then just use the normal input and zeros 313 if h_t.sum() == 0: 314 return args, kwargs 315 # If its not the first one then return the input it got with its own 316 # h_t and c_t to replace neurons 317 else: 318 return (args[0], (self.h_t_d, self.c_t_d)), kwargs
Filter LSTMCell input for dendrite based on initialization state.
Checks if this is the first time step (all zeros in h_t) or a subsequent step. For the first step, passes through the original inputs. For subsequent steps, replaces the neuron's hidden state with the dendrite's own internal state from the previous iteration.
Parameters
- *args (tuple): Contains (input, (h_t, c_t)) where input is the external input and (h_t, c_t) is the neuron's recurrent state.
- **kwargs (dict): Keyword arguments to pass through.
Returns
- tuple: ((processed_input, processed_state), kwargs) for dendrite call.
def
post_d(self, *args, **kwargs):
320 def post_d(self, *args, **kwargs): 321 """ 322 Extract and store dendrite's LSTM state for next iteration. 323 324 Separates the dendrite's hidden and cell states from its output tuple, 325 stores both for use in the next time step, and returns only the hidden 326 state to be combined with the neuron's output. 327 328 Parameters 329 ---------- 330 *args : tuple 331 Contains dendrite LSTM output tuple (h_t, c_t). 332 **kwargs : dict 333 Unused keyword arguments. 334 335 Returns 336 ------- 337 torch.Tensor 338 Hidden state h_t to be added to the neuron output. 339 """ 340 h_t = args[0][0] 341 c_t = args[0][1] 342 self.h_t_d = h_t 343 self.c_t_d = c_t 344 return h_t
Extract and store dendrite's LSTM state for next iteration.
Separates the dendrite's hidden and cell states from its output tuple, stores both for use in the next time step, and returns only the hidden state to be combined with the neuron's output.
Parameters
- *args (tuple): Contains dendrite LSTM output tuple (h_t, c_t).
- **kwargs (dict): Unused keyword arguments.
Returns
- torch.Tensor: Hidden state h_t to be added to the neuron output.
def
clear_processor(self):
346 def clear_processor(self): 347 """ 348 Clear all stored LSTM states. 349 350 Removes dendrite hidden state (h_t_d), dendrite cell state (c_t_d), 351 and temporarily stored neuron cell state (c_t_n). Safe to call even 352 if attributes don't exist. 353 """ 354 if hasattr(self, "h_t_d"): 355 delattr(self, "h_t_d") 356 if hasattr(self, "c_t_d"): 357 delattr(self, "c_t_d") 358 if hasattr(self, "c_t_n"): 359 delattr(self, "c_t_n")
Clear all stored LSTM states.
Removes dendrite hidden state (h_t_d), dendrite cell state (c_t_d), and temporarily stored neuron cell state (c_t_n). Safe to call even if attributes don't exist.
363class LSTMProcessor(PAIProcessor): 364 """Processor for LSTM to handle hidden and output states.""" 365 366 def post_n1(self, *args, **kwargs): 367 """ 368 Extract hidden state from LSTM output for dendrite processing. 369 370 Separates the hidden state from the output in the 371 LSTM output tuple. Stores the hidden state temporarily since only the 372 output state should be modified by dendrites. 373 374 Parameters 375 ---------- 376 *args : tuple 377 Contains LSTM output tuple (output, hidden) as first element. 378 **kwargs : dict 379 Unused keyword arguments. 380 381 Returns 382 ------- 383 torch.Tensor 384 Output state to be passed to dendrite processing. 385 """ 386 output = args[0][0] 387 hidden = args[0][1] 388 # Store the hidden state temporarily and just use the output state 389 # to do Dendrite functions 390 self.hidden_n = hidden 391 return output 392 393 def post_n2(self, *args, **kwargs): 394 """ 395 Recombine dendrite-modified output with hidden tuple. 396 397 Takes the output state that has been modified by dendrite operations 398 and combines it with the stored hidden state to produce the complete 399 LSTM output tuple. 400 401 Parameters 402 ---------- 403 *args : tuple 404 Contains the dendrite-modified output state. 405 **kwargs : dict 406 Unused keyword arguments. 407 408 Returns 409 ------- 410 tuple 411 Complete LSTM output (output, hidden) where output has been modified. 412 """ 413 output = args[0] 414 return output, self.hidden_n 415 416 def pre_d(self, *args, **kwargs): 417 """ 418 LSTM input is just the tensor which also goes to the dendrite 419 420 Parameters 421 ---------- 422 *args : 423 Input tensor 424 **kwargs : dict 425 Empty 426 427 Returns 428 ------- 429 tuple 430 (output, hidden) 431 """ 432 return args, kwargs 433 434 def post_d(self, *args, **kwargs): 435 """ 436 Extract dendrite's output to combine. 437 438 Parameters 439 ---------- 440 *args : tuple 441 Contains dendrite LSTM output tuple (output, hidden). 442 **kwargs : dict 443 Unused keyword arguments. 444 445 Returns 446 ------- 447 torch.Tensor 448 Output state to be added to the neuron output. 449 """ 450 output = args[0][0] 451 hidden = args[0][1] 452 return output 453 454 def clear_processor(self): 455 """ 456 Clear all stored LSTM states. 457 458 """ 459 if hasattr(self, "hidden_n"): 460 delattr(self, "hidden_n")
Processor for LSTM to handle hidden and output states.
def
post_n1(self, *args, **kwargs):
366 def post_n1(self, *args, **kwargs): 367 """ 368 Extract hidden state from LSTM output for dendrite processing. 369 370 Separates the hidden state from the output in the 371 LSTM output tuple. Stores the hidden state temporarily since only the 372 output state should be modified by dendrites. 373 374 Parameters 375 ---------- 376 *args : tuple 377 Contains LSTM output tuple (output, hidden) as first element. 378 **kwargs : dict 379 Unused keyword arguments. 380 381 Returns 382 ------- 383 torch.Tensor 384 Output state to be passed to dendrite processing. 385 """ 386 output = args[0][0] 387 hidden = args[0][1] 388 # Store the hidden state temporarily and just use the output state 389 # to do Dendrite functions 390 self.hidden_n = hidden 391 return output
Extract hidden state from LSTM output for dendrite processing.
Separates the hidden state from the output in the LSTM output tuple. Stores the hidden state temporarily since only the output state should be modified by dendrites.
Parameters
- *args (tuple): Contains LSTM output tuple (output, hidden) as first element.
- **kwargs (dict): Unused keyword arguments.
Returns
- torch.Tensor: Output state to be passed to dendrite processing.
def
post_n2(self, *args, **kwargs):
393 def post_n2(self, *args, **kwargs): 394 """ 395 Recombine dendrite-modified output with hidden tuple. 396 397 Takes the output state that has been modified by dendrite operations 398 and combines it with the stored hidden state to produce the complete 399 LSTM output tuple. 400 401 Parameters 402 ---------- 403 *args : tuple 404 Contains the dendrite-modified output state. 405 **kwargs : dict 406 Unused keyword arguments. 407 408 Returns 409 ------- 410 tuple 411 Complete LSTM output (output, hidden) where output has been modified. 412 """ 413 output = args[0] 414 return output, self.hidden_n
Recombine dendrite-modified output with hidden tuple.
Takes the output state that has been modified by dendrite operations and combines it with the stored hidden state to produce the complete LSTM output tuple.
Parameters
- *args (tuple): Contains the dendrite-modified output state.
- **kwargs (dict): Unused keyword arguments.
Returns
- tuple: Complete LSTM output (output, hidden) where output has been modified.
def
pre_d(self, *args, **kwargs):
416 def pre_d(self, *args, **kwargs): 417 """ 418 LSTM input is just the tensor which also goes to the dendrite 419 420 Parameters 421 ---------- 422 *args : 423 Input tensor 424 **kwargs : dict 425 Empty 426 427 Returns 428 ------- 429 tuple 430 (output, hidden) 431 """ 432 return args, kwargs
LSTM input is just the tensor which also goes to the dendrite
Parameters
- *args (): Input tensor
- **kwargs (dict): Empty
Returns
- tuple: (output, hidden)
def
post_d(self, *args, **kwargs):
434 def post_d(self, *args, **kwargs): 435 """ 436 Extract dendrite's output to combine. 437 438 Parameters 439 ---------- 440 *args : tuple 441 Contains dendrite LSTM output tuple (output, hidden). 442 **kwargs : dict 443 Unused keyword arguments. 444 445 Returns 446 ------- 447 torch.Tensor 448 Output state to be added to the neuron output. 449 """ 450 output = args[0][0] 451 hidden = args[0][1] 452 return output
Extract dendrite's output to combine.
Parameters
- *args (tuple): Contains dendrite LSTM output tuple (output, hidden).
- **kwargs (dict): Unused keyword arguments.
Returns
- torch.Tensor: Output state to be added to the neuron output.
463class LSTMProcessorLastHidden(PAIProcessor): 464 """Processor for LSTM to forward the last hidden.""" 465 466 def post_n1(self, *args, **kwargs): 467 """ 468 Extract the last hidden to combine with dendrites 469 470 Parameters 471 ---------- 472 *args : tuple 473 Contains LSTM output tuple (output, hidden) as first element. 474 **kwargs : dict 475 Unused keyword arguments. 476 477 Returns 478 ------- 479 torch.Tensor 480 Output state to be passed to dendrite processing. 481 """ 482 ignored_output = args[0][0] 483 last_hidden = args[0][1][-1] 484 485 return last_hidden 486 487 def post_n2(self, *args, **kwargs): 488 """ 489 Recombine dendrite-modified last hidden, and append None just to maintain output format 490 491 Parameters 492 ---------- 493 *args : tuple 494 Contains the dendrite-modified output state. 495 **kwargs : dict 496 Unused keyword arguments. 497 498 Returns 499 ------- 500 tuple 501 Complete LSTM output (output, hidden) where output has been modified. 502 """ 503 combined_last_hidden = args[0] 504 return None, combined_last_hidden 505 506 def pre_d(self, *args, **kwargs): 507 """ 508 LSTM input is just the tensor which also goes to the dendrite 509 510 Parameters 511 ---------- 512 *args : 513 Input tensor 514 **kwargs : dict 515 Empty 516 517 Returns 518 ------- 519 tuple 520 (output, hidden) 521 """ 522 return args, kwargs 523 524 def post_d(self, *args, **kwargs): 525 """ 526 Extract extract the dendrites last hidden to combine with neurons. 527 528 Parameters 529 ---------- 530 *args : tuple 531 Contains dendrite LSTM output tuple (output, hidden). 532 **kwargs : dict 533 Unused keyword arguments. 534 535 Returns 536 ------- 537 torch.Tensor 538 Output state to be added to the neuron output. 539 """ 540 ignored_output = args[0][0] 541 last_hidden = args[0][1][-1] 542 return last_hidden 543 544 def clear_processor(self): 545 # Nothing is stored 546 pass
Processor for LSTM to forward the last hidden.
def
post_n1(self, *args, **kwargs):
466 def post_n1(self, *args, **kwargs): 467 """ 468 Extract the last hidden to combine with dendrites 469 470 Parameters 471 ---------- 472 *args : tuple 473 Contains LSTM output tuple (output, hidden) as first element. 474 **kwargs : dict 475 Unused keyword arguments. 476 477 Returns 478 ------- 479 torch.Tensor 480 Output state to be passed to dendrite processing. 481 """ 482 ignored_output = args[0][0] 483 last_hidden = args[0][1][-1] 484 485 return last_hidden
Extract the last hidden to combine with dendrites
Parameters
- *args (tuple): Contains LSTM output tuple (output, hidden) as first element.
- **kwargs (dict): Unused keyword arguments.
Returns
- torch.Tensor: Output state to be passed to dendrite processing.
def
post_n2(self, *args, **kwargs):
487 def post_n2(self, *args, **kwargs): 488 """ 489 Recombine dendrite-modified last hidden, and append None just to maintain output format 490 491 Parameters 492 ---------- 493 *args : tuple 494 Contains the dendrite-modified output state. 495 **kwargs : dict 496 Unused keyword arguments. 497 498 Returns 499 ------- 500 tuple 501 Complete LSTM output (output, hidden) where output has been modified. 502 """ 503 combined_last_hidden = args[0] 504 return None, combined_last_hidden
Recombine dendrite-modified last hidden, and append None just to maintain output format
Parameters
- *args (tuple): Contains the dendrite-modified output state.
- **kwargs (dict): Unused keyword arguments.
Returns
- tuple: Complete LSTM output (output, hidden) where output has been modified.
def
pre_d(self, *args, **kwargs):
506 def pre_d(self, *args, **kwargs): 507 """ 508 LSTM input is just the tensor which also goes to the dendrite 509 510 Parameters 511 ---------- 512 *args : 513 Input tensor 514 **kwargs : dict 515 Empty 516 517 Returns 518 ------- 519 tuple 520 (output, hidden) 521 """ 522 return args, kwargs
LSTM input is just the tensor which also goes to the dendrite
Parameters
- *args (): Input tensor
- **kwargs (dict): Empty
Returns
- tuple: (output, hidden)
def
post_d(self, *args, **kwargs):
524 def post_d(self, *args, **kwargs): 525 """ 526 Extract extract the dendrites last hidden to combine with neurons. 527 528 Parameters 529 ---------- 530 *args : tuple 531 Contains dendrite LSTM output tuple (output, hidden). 532 **kwargs : dict 533 Unused keyword arguments. 534 535 Returns 536 ------- 537 torch.Tensor 538 Output state to be added to the neuron output. 539 """ 540 ignored_output = args[0][0] 541 last_hidden = args[0][1][-1] 542 return last_hidden
Extract extract the dendrites last hidden to combine with neurons.
Parameters
- *args (tuple): Contains dendrite LSTM output tuple (output, hidden).
- **kwargs (dict): Unused keyword arguments.
Returns
- torch.Tensor: Output state to be added to the neuron output.
class
ResNetPAI(torch.nn.modules.module.Module):
548class ResNetPAI(nn.Module): 549 """PB-compatible ResNet wrapper. 550 551 All normalization layers should be wrapped in a PAISequential, or other 552 wrapped module. When working with a predefined model the following shows 553 an example of how to create a module for modules_to_replace. 554 """ 555 556 def __init__(self, other_resnet): 557 """Initialize ResNetPAI from existing ResNet model. 558 559 Parameters 560 ---------- 561 *args : other_resnet : torchvision.models.resnet.ResNet 562 An existing ResNet model to convert to PAI-compatible format. 563 """ 564 super(ResNetPAI, self).__init__() 565 566 # For the most part, just copy the exact values from the original module 567 self._norm_layer = other_resnet._norm_layer 568 self.inplanes = other_resnet.inplanes 569 self.dilation = other_resnet.dilation 570 self.groups = other_resnet.groups 571 self.base_width = other_resnet.base_width 572 573 # For the component to be changed, define a PAISequential with the old 574 # modules included 575 self.b1 = GPA.PAISequential([other_resnet.conv1, other_resnet.bn1]) 576 577 self.relu = other_resnet.relu 578 self.maxpool = other_resnet.maxpool 579 580 for i in range(1, 5): 581 layer_name = "layer" + str(i) 582 original_layer = getattr(other_resnet, layer_name) 583 pb_layer = self._make_layer_pb(original_layer, other_resnet, i) 584 setattr(self, layer_name, pb_layer) 585 586 self.avgpool = other_resnet.avgpool 587 self.fc = other_resnet.fc 588 589 def _make_layer_pb(self, other_block_set, other_resnet, block_id): 590 """Convert ResNet layer blocks to PB-compatible format. 591 592 Parameters 593 ---------- 594 other_block_set : torch.vision.models.resnet.any_block 595 A set of blocks from the original ResNet model. 596 other_resnet : torchvision.models.resnet.ResNet 597 The original ResNet model. 598 block_id : int 599 The layer number being converted. 600 Returns 601 ------- 602 nn.Sequential 603 A sequential container with the converted blocks. 604 """ 605 layers = [] 606 for i in range(len(other_block_set)): 607 block_type = type(other_block_set[i]) 608 if block_type == resnet_pt.BasicBlock: 609 layers.append(other_block_set[i]) 610 elif block_type == resnet_pt.Bottleneck: 611 layers.append(other_block_set[i]) 612 else: 613 print( 614 "Your resnet uses a block type that has not been " 615 "accounted for. Customization might be required." 616 ) 617 layer_name = "layer" + str(block_id) 618 print(type(getattr(other_resnet, layer_name))) 619 pdb.set_trace() 620 return nn.Sequential(*layers) 621 622 def _forward_impl(self, x): 623 """Implementation of the forward pass. 624 625 Parameters 626 ---------- 627 x : torch.Tensor 628 Input tensor to the network. 629 630 Returns 631 ------- 632 torch.Tensor 633 Output tensor from the network. 634 """ 635 # Modified b1 rather than conv1 and bn1 636 x = self.b1(x) 637 # Rest of forward remains the same 638 x = F.relu(x) 639 x = self.maxpool(x) 640 641 x = self.layer1(x) 642 x = self.layer2(x) 643 x = self.layer3(x) 644 x = self.layer4(x) 645 646 x = self.avgpool(x) 647 x = torch.flatten(x, 1) 648 x = self.fc(x) 649 650 return x 651 652 def forward(self, x): 653 """Forward pass through the network. 654 655 Parameters 656 ---------- 657 x : torch.Tensor 658 Input tensor to the network. 659 660 Returns 661 ------- 662 torch.Tensor 663 Output tensor from the network. 664 """ 665 return self._forward_impl(x)
PB-compatible ResNet wrapper.
All normalization layers should be wrapped in a PAISequential, or other wrapped module. When working with a predefined model the following shows an example of how to create a module for modules_to_replace.
ResNetPAI(other_resnet)
556 def __init__(self, other_resnet): 557 """Initialize ResNetPAI from existing ResNet model. 558 559 Parameters 560 ---------- 561 *args : other_resnet : torchvision.models.resnet.ResNet 562 An existing ResNet model to convert to PAI-compatible format. 563 """ 564 super(ResNetPAI, self).__init__() 565 566 # For the most part, just copy the exact values from the original module 567 self._norm_layer = other_resnet._norm_layer 568 self.inplanes = other_resnet.inplanes 569 self.dilation = other_resnet.dilation 570 self.groups = other_resnet.groups 571 self.base_width = other_resnet.base_width 572 573 # For the component to be changed, define a PAISequential with the old 574 # modules included 575 self.b1 = GPA.PAISequential([other_resnet.conv1, other_resnet.bn1]) 576 577 self.relu = other_resnet.relu 578 self.maxpool = other_resnet.maxpool 579 580 for i in range(1, 5): 581 layer_name = "layer" + str(i) 582 original_layer = getattr(other_resnet, layer_name) 583 pb_layer = self._make_layer_pb(original_layer, other_resnet, i) 584 setattr(self, layer_name, pb_layer) 585 586 self.avgpool = other_resnet.avgpool 587 self.fc = other_resnet.fc
Initialize ResNetPAI from existing ResNet model.
Parameters
- *args : other_resnet (torchvision.models.resnet.ResNet): An existing ResNet model to convert to PAI-compatible format.
def
forward(self, x):
652 def forward(self, x): 653 """Forward pass through the network. 654 655 Parameters 656 ---------- 657 x : torch.Tensor 658 Input tensor to the network. 659 660 Returns 661 ------- 662 torch.Tensor 663 Output tensor from the network. 664 """ 665 return self._forward_impl(x)
Forward pass through the network.
Parameters
- x (torch.Tensor): Input tensor to the network.
Returns
- torch.Tensor: Output tensor from the network.
class
ResNetPAIPreFC(torch.nn.modules.module.Module):
668class ResNetPAIPreFC(nn.Module): 669 """PB-compatible ResNet wrapper. 670 671 All normalization layers should be wrapped in a PAISequential, or other 672 wrapped module. When working with a predefined model the following shows 673 an example of how to create a module for modules_to_replace. 674 """ 675 676 def __init__(self, other_resnet): 677 """Initialize ResNetPAI from existing ResNet model. 678 679 Parameters 680 ---------- 681 *args : other_resnet : torchvision.models.resnet.ResNet 682 An existing ResNet model to convert to PAI-compatible format. 683 """ 684 super(ResNetPAIPreFC, self).__init__() 685 686 # For the most part, just copy the exact values from the original module 687 self._norm_layer = other_resnet._norm_layer 688 self.inplanes = other_resnet.inplanes 689 self.dilation = other_resnet.dilation 690 self.groups = other_resnet.groups 691 self.base_width = other_resnet.base_width 692 693 # For the component to be changed, define a PAISequential with the old 694 # modules included 695 self.conv1 = other_resnet.conv1 696 self.bn1 = other_resnet.bn1 697 698 self.relu = other_resnet.relu 699 self.maxpool = other_resnet.maxpool 700 701 for i in range(1, 5): 702 layer_name = "layer" + str(i) 703 original_layer = getattr(other_resnet, layer_name) 704 pb_layer = self._make_layer_pb(original_layer, other_resnet, i) 705 setattr(self, layer_name, pb_layer) 706 707 self.avgpool = other_resnet.avgpool 708 709 # Create pre_fc layer with dimensions matching layer4 output (same as fc input) 710 fc_in_features = other_resnet.fc.in_features 711 self.pre_fc = nn.Linear(fc_in_features, fc_in_features) 712 713 self.fc = other_resnet.fc 714 715 def _make_layer_pb(self, other_block_set, other_resnet, block_id): 716 """Convert ResNet layer blocks to PB-compatible format. 717 718 Parameters 719 ---------- 720 other_block_set : torch.vision.models.resnet.any_block 721 A set of blocks from the original ResNet model. 722 other_resnet : torchvision.models.resnet.ResNet 723 The original ResNet model. 724 block_id : int 725 The layer number being converted. 726 Returns 727 ------- 728 nn.Sequential 729 A sequential container with the converted blocks. 730 """ 731 layers = [] 732 for i in range(len(other_block_set)): 733 block_type = type(other_block_set[i]) 734 if block_type == resnet_pt.BasicBlock: 735 layers.append(other_block_set[i]) 736 elif block_type == resnet_pt.Bottleneck: 737 layers.append(other_block_set[i]) 738 else: 739 print( 740 "Your resnet uses a block type that has not been " 741 "accounted for. Customization might be required." 742 ) 743 layer_name = "layer" + str(block_id) 744 print(type(getattr(other_resnet, layer_name))) 745 pdb.set_trace() 746 return nn.Sequential(*layers) 747 748 def _forward_impl(self, x): 749 """Implementation of the forward pass. 750 751 Parameters 752 ---------- 753 x : torch.Tensor 754 Input tensor to the network. 755 756 Returns 757 ------- 758 torch.Tensor 759 Output tensor from the network. 760 """ 761 # Modified b1 rather than conv1 and bn1 762 x = self.conv1(x) 763 x = self.bn1(x) 764 # Rest of forward remains the same 765 x = F.relu(x) 766 x = self.maxpool(x) 767 768 x = self.layer1(x) 769 x = self.layer2(x) 770 x = self.layer3(x) 771 x = self.layer4(x) 772 773 x = self.avgpool(x) 774 x = torch.flatten(x, 1) 775 x = self.pre_fc(x) 776 x = F.relu(x) 777 x = self.fc(x) 778 779 return x 780 781 def forward(self, x): 782 """Forward pass through the network. 783 784 Parameters 785 ---------- 786 x : torch.Tensor 787 Input tensor to the network. 788 789 Returns 790 ------- 791 torch.Tensor 792 Output tensor from the network. 793 """ 794 return self._forward_impl(x)
PB-compatible ResNet wrapper.
All normalization layers should be wrapped in a PAISequential, or other wrapped module. When working with a predefined model the following shows an example of how to create a module for modules_to_replace.
ResNetPAIPreFC(other_resnet)
676 def __init__(self, other_resnet): 677 """Initialize ResNetPAI from existing ResNet model. 678 679 Parameters 680 ---------- 681 *args : other_resnet : torchvision.models.resnet.ResNet 682 An existing ResNet model to convert to PAI-compatible format. 683 """ 684 super(ResNetPAIPreFC, self).__init__() 685 686 # For the most part, just copy the exact values from the original module 687 self._norm_layer = other_resnet._norm_layer 688 self.inplanes = other_resnet.inplanes 689 self.dilation = other_resnet.dilation 690 self.groups = other_resnet.groups 691 self.base_width = other_resnet.base_width 692 693 # For the component to be changed, define a PAISequential with the old 694 # modules included 695 self.conv1 = other_resnet.conv1 696 self.bn1 = other_resnet.bn1 697 698 self.relu = other_resnet.relu 699 self.maxpool = other_resnet.maxpool 700 701 for i in range(1, 5): 702 layer_name = "layer" + str(i) 703 original_layer = getattr(other_resnet, layer_name) 704 pb_layer = self._make_layer_pb(original_layer, other_resnet, i) 705 setattr(self, layer_name, pb_layer) 706 707 self.avgpool = other_resnet.avgpool 708 709 # Create pre_fc layer with dimensions matching layer4 output (same as fc input) 710 fc_in_features = other_resnet.fc.in_features 711 self.pre_fc = nn.Linear(fc_in_features, fc_in_features) 712 713 self.fc = other_resnet.fc
Initialize ResNetPAI from existing ResNet model.
Parameters
- *args : other_resnet (torchvision.models.resnet.ResNet): An existing ResNet model to convert to PAI-compatible format.
def
forward(self, x):
781 def forward(self, x): 782 """Forward pass through the network. 783 784 Parameters 785 ---------- 786 x : torch.Tensor 787 Input tensor to the network. 788 789 Returns 790 ------- 791 torch.Tensor 792 Output tensor from the network. 793 """ 794 return self._forward_impl(x)
Forward pass through the network.
Parameters
- x (torch.Tensor): Input tensor to the network.
Returns
- torch.Tensor: Output tensor from the network.