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        Parameters
145        ----------
146        None
147
148        Returns
149        -------
150        None
151            This function does not return a value.
152        """
153        pass
154
155
156# General multi output processor for any number that ignores later ones
157class MultiOutputProcessor:
158    """Processor for handling multiple outputs, ignoring later ones."""
159
160    def post_n1(self, *args, **kwargs):
161        """Saves extra outputs and returns the first output.
162
163        Parameters
164        ----------
165        *args : tuple
166            Contains the modules output tuple.
167        **kwargs : dict
168            Unused keyword arguments.
169
170        Returns
171        -------
172        torch.Tensor
173            The first tensor of the tuple
174        """
175        out = args[0][0]
176        extra_out = args[0][1:]
177        self.extra_out = extra_out
178        return out
179
180    def post_n2(self, *args, **kwargs):
181        """Combine output with stored extra outputs.
182
183        Parameters
184        ----------
185        *args : torch.tensor
186            The first tensor combined with dendrite output.
187        **kwargs : dict
188            Unused keyword arguments.
189
190        Returns
191        -------
192        tuple
193            The recombined output tuple wth the new first output modified
194        """
195        out = args[0]
196        if isinstance(self.extra_out, tuple):
197            return (out,) + self.extra_out
198        else:
199            return (out,) + (self.extra_out,)
200
201    def pre_d(self, *args, **kwargs):
202        """Pass through arguments unchanged for dendrite preprocessing.
203
204        Parameters
205        ----------
206        *args : tuple
207            Positional arguments containing inputs to the PAI module.
208        **kwargs : dict
209            Keyword arguments containing inputs to the PAI module.
210
211        Returns
212        -------
213        args : tuple
214            Positional arguments containing inputs to the PAI module.
215        kwargs : dict
216            Keyword arguments containing inputs to the PAI module.
217        """
218        return args, kwargs
219
220    def post_d(self, *args, **kwargs):
221        """Extract first output for dendrite postprocessing.
222
223        Parameters
224        ----------
225        *args : tuple
226            Contains the dendrite modules output tuple.
227        **kwargs : dict
228            Unused keyword arguments.
229
230        Returns
231        -------
232        torch.Tensor
233            The first tensor of the tuple
234        """
235        out = args[0][0]
236        return out
237
238    def clear_processor(self):
239        """Clear stored processor state.
240
241        Parameters
242        ----------
243        None
244
245        Returns
246        -------
247        None
248            This function does not return a value.
249        """
250
251        if hasattr(self, "extra_out"):
252            delattr(self, "extra_out")
253
254class LSTMCellProcessor(PAIProcessor):
255    """Processor for LSTM cells to handle hidden and cell states."""
256
257    def post_n1(self, *args, **kwargs):
258        """
259        Extract hidden state from LSTM output for dendrite processing.
260
261        Separates the hidden state (h_t) from the cell state (c_t) in the
262        LSTM output tuple. Stores the cell state temporarily since only the
263        hidden state should be modified by dendrites.
264
265        Parameters
266        ----------
267        *args : tuple
268            Contains LSTM output tuple (h_t, c_t) as first element.
269        **kwargs : dict
270            Unused keyword arguments.
271
272        Returns
273        -------
274        torch.Tensor
275            Hidden state h_t to be passed to dendrite processing.
276        """
277        h_t = args[0][0]
278        c_t = args[0][1]
279        # Store the cell state temporarily and just use the hidden state
280        # to do Dendrite functions
281        self.c_t_n = c_t
282        return h_t
283
284    def post_n2(self, *args, **kwargs):
285        """
286        Recombine dendrite-modified hidden state with cell state.
287
288        Takes the hidden state that has been modified by dendrite operations
289        and combines it with the stored cell state to produce the complete
290        LSTM output tuple.
291
292        Parameters
293        ----------
294        *args : tuple
295            Contains the dendrite-modified hidden state h_t.
296        **kwargs : dict
297            Unused keyword arguments.
298
299        Returns
300        -------
301        tuple
302            Complete LSTM output (h_t, c_t) where h_t has been modified.
303        """
304        h_t = args[0]
305        return h_t, self.c_t_n
306
307    def pre_d(self, *args, **kwargs):
308        """
309        Filter LSTMCell input for dendrite based on initialization state.
310
311        Checks if this is the first time step (all zeros in h_t) or a
312        subsequent step. For the first step, passes through the original
313        inputs. For subsequent steps, replaces the neuron's hidden state
314        with the dendrite's own internal state from the previous iteration.
315
316        Parameters
317        ----------
318        *args : tuple
319            Contains (input, (h_t, c_t)) where input is the external input
320            and (h_t, c_t) is the neuron's recurrent state.
321        **kwargs : dict
322            Keyword arguments to pass through.
323
324        Returns
325        -------
326        tuple
327            ((processed_input, processed_state), kwargs) for dendrite call.
328        """
329        h_t = args[1][0]
330        # If its the initial step then just use the normal input and zeros
331        if h_t.sum() == 0:
332            return args, kwargs
333        # If its not the first one then return the input it got with its own
334        # h_t and c_t to replace neurons
335        else:
336            return (args[0], (self.h_t_d, self.c_t_d)), kwargs
337
338    def post_d(self, *args, **kwargs):
339        """
340        Extract and store dendrite's LSTM state for next iteration.
341
342        Separates the dendrite's hidden and cell states from its output tuple,
343        stores both for use in the next time step, and returns only the hidden
344        state to be combined with the neuron's output.
345
346        Parameters
347        ----------
348        *args : tuple
349            Contains dendrite LSTM output tuple (h_t, c_t).
350        **kwargs : dict
351            Unused keyword arguments.
352
353        Returns
354        -------
355        torch.Tensor
356            Hidden state h_t to be added to the neuron output.
357        """
358        h_t = args[0][0]
359        c_t = args[0][1]
360        self.h_t_d = h_t
361        self.c_t_d = c_t
362        return h_t
363
364    def clear_processor(self):
365        """
366        Clear all stored LSTM states.
367
368        Removes dendrite hidden state (h_t_d), dendrite cell state (c_t_d),
369        and temporarily stored neuron cell state (c_t_n). Safe to call even
370        if attributes don't exist.
371
372        Parameters
373        ----------
374        None
375
376        Returns
377        -------
378        None
379            This function does not return a value.
380        """
381        if hasattr(self, "h_t_d"):
382            delattr(self, "h_t_d")
383        if hasattr(self, "c_t_d"):
384            delattr(self, "c_t_d")
385        if hasattr(self, "c_t_n"):
386            delattr(self, "c_t_n")
387
388
389
390class LSTMProcessor(PAIProcessor):
391    """Processor for LSTM to handle hidden and output states."""
392
393    def post_n1(self, *args, **kwargs):
394        """
395        Extract hidden state from LSTM output for dendrite processing.
396
397        Separates the hidden state from the output in the
398        LSTM output tuple. Stores the hidden state temporarily since only the
399        output state should be modified by dendrites.
400
401        Parameters
402        ----------
403        *args : tuple
404            Contains LSTM output tuple (output, hidden) as first element.
405        **kwargs : dict
406            Unused keyword arguments.
407
408        Returns
409        -------
410        torch.Tensor
411            Output state to be passed to dendrite processing.
412        """
413        output = args[0][0]
414        hidden = args[0][1]
415        # Store the hidden state temporarily and just use the output state
416        # to do Dendrite functions
417        self.hidden_n = hidden
418        return output
419
420    def post_n2(self, *args, **kwargs):
421        """
422        Recombine dendrite-modified output with hidden tuple.
423
424        Takes the output state that has been modified by dendrite operations
425        and combines it with the stored hidden state to produce the complete
426        LSTM output tuple.
427
428        Parameters
429        ----------
430        *args : tuple
431            Contains the dendrite-modified output state.
432        **kwargs : dict
433            Unused keyword arguments.
434
435        Returns
436        -------
437        tuple
438            Complete LSTM output (output, hidden) where output has been modified.
439        """
440        output = args[0]
441        return output, self.hidden_n
442
443    def pre_d(self, *args, **kwargs):
444        """
445        LSTM input is just the tensor which also goes to the dendrite
446
447        Parameters
448        ----------
449        *args : 
450            Input tensor
451        **kwargs : dict
452            Empty
453
454        Returns
455        -------
456        tuple
457            (output, hidden)
458        """
459        return args, kwargs
460        
461    def post_d(self, *args, **kwargs):
462        """
463        Extract dendrite's output to combine.
464
465        Parameters
466        ----------
467        *args : tuple
468            Contains dendrite LSTM output tuple (output, hidden).
469        **kwargs : dict
470            Unused keyword arguments.
471
472        Returns
473        -------
474        torch.Tensor
475            Output state to be added to the neuron output.
476        """
477        output = args[0][0]
478        hidden = args[0][1]
479        return output
480
481    def clear_processor(self):
482        """
483        Clear all stored LSTM states.
484
485
486        Parameters
487        ----------
488        None
489
490        Returns
491        -------
492        None
493            This function does not return a value.
494        """
495        if hasattr(self, "hidden_n"):
496            delattr(self, "hidden_n")
497
498
499class LSTMProcessorLastHidden(PAIProcessor):
500    """Processor for LSTM to forward the last hidden."""
501
502    def post_n1(self, *args, **kwargs):
503        """
504        Extract the last hidden to combine with dendrites
505
506        Parameters
507        ----------
508        *args : tuple
509            Contains LSTM output tuple (output, hidden) as first element.
510        **kwargs : dict
511            Unused keyword arguments.
512
513        Returns
514        -------
515        torch.Tensor
516            Output state to be passed to dendrite processing.
517        """
518        ignored_output = args[0][0]
519        last_hidden = args[0][1][-1]
520
521        return last_hidden
522
523    def post_n2(self, *args, **kwargs):
524        """
525        Recombine dendrite-modified last hidden, and append None just to maintain output format
526
527        Parameters
528        ----------
529        *args : tuple
530            Contains the dendrite-modified output state.
531        **kwargs : dict
532            Unused keyword arguments.
533
534        Returns
535        -------
536        tuple
537            Complete LSTM output (output, hidden) where output has been modified.
538        """
539        combined_last_hidden = args[0]
540        return None, combined_last_hidden
541
542    def pre_d(self, *args, **kwargs):
543        """
544        LSTM input is just the tensor which also goes to the dendrite
545
546        Parameters
547        ----------
548        *args : 
549            Input tensor
550        **kwargs : dict
551            Empty
552
553        Returns
554        -------
555        tuple
556            (output, hidden)
557        """
558        return args, kwargs
559        
560    def post_d(self, *args, **kwargs):
561        """
562        Extract extract the dendrites last hidden to combine with neurons.
563
564        Parameters
565        ----------
566        *args : tuple
567            Contains dendrite LSTM output tuple (output, hidden).
568        **kwargs : dict
569            Unused keyword arguments.
570
571        Returns
572        -------
573        torch.Tensor
574            Output state to be added to the neuron output.
575        """
576        ignored_output = args[0][0]
577        last_hidden = args[0][1][-1]
578        return last_hidden
579
580    def clear_processor(self):
581        """Clear processor state.
582
583        Notes
584        -----
585        This processor keeps no persistent internal state.
586
587        Parameters
588        ----------
589        None
590
591        Returns
592        -------
593        None
594            This function does not return a value.
595        """
596        # Nothing is stored
597        pass
598
599class ResNetPAI(nn.Module):
600    """PB-compatible ResNet wrapper.
601
602    All normalization layers should be wrapped in a PAISequential, or other
603    wrapped module. When working with a predefined model the following shows
604    an example of how to create a module for modules_to_replace.
605    """
606
607    def __init__(self, other_resnet):
608        """Initialize ResNetPAI from existing ResNet model.
609
610        Parameters
611        ----------
612        *args : other_resnet : torchvision.models.resnet.ResNet
613            An existing ResNet model to convert to PAI-compatible format.
614        """
615        super(ResNetPAI, self).__init__()
616
617        # For the most part, just copy the exact values from the original module
618        self._norm_layer = other_resnet._norm_layer
619        self.inplanes = other_resnet.inplanes
620        self.dilation = other_resnet.dilation
621        self.groups = other_resnet.groups
622        self.base_width = other_resnet.base_width
623
624        # For the component to be changed, define a PAISequential with the old
625        # modules included
626        self.b1 = GPA.PAISequential([other_resnet.conv1, other_resnet.bn1])
627
628        self.relu = other_resnet.relu
629        self.maxpool = other_resnet.maxpool
630
631        for i in range(1, 5):
632            layer_name = "layer" + str(i)
633            original_layer = getattr(other_resnet, layer_name)
634            pb_layer = self._make_layer_pb(original_layer, other_resnet, i)
635            setattr(self, layer_name, pb_layer)
636
637        self.avgpool = other_resnet.avgpool
638        self.fc = other_resnet.fc
639
640    def _make_layer_pb(self, other_block_set, other_resnet, block_id):
641        """Convert ResNet layer blocks to PB-compatible format.
642
643        Parameters
644        ----------
645        other_block_set : torch.vision.models.resnet.any_block
646            A set of blocks from the original ResNet model.
647        other_resnet : torchvision.models.resnet.ResNet
648            The original ResNet model.
649        block_id : int
650            The layer number being converted.
651        Returns
652        -------
653        nn.Sequential
654            A sequential container with the converted blocks.
655        """
656        layers = []
657        for i in range(len(other_block_set)):
658            block_type = type(other_block_set[i])
659            if block_type == resnet_pt.BasicBlock:
660                layers.append(other_block_set[i])
661            elif block_type == resnet_pt.Bottleneck:
662                layers.append(other_block_set[i])
663            else:
664                print(
665                    "Your resnet uses a block type that has not been "
666                    "accounted for. Customization might be required."
667                )
668                layer_name = "layer" + str(block_id)
669                print(type(getattr(other_resnet, layer_name)))
670                pdb.set_trace()
671        return nn.Sequential(*layers)
672
673    def _forward_impl(self, x):
674        """Implementation of the forward pass.
675
676        Parameters
677        ----------
678        x : torch.Tensor
679            Input tensor to the network.
680
681        Returns
682        -------
683        torch.Tensor
684            Output tensor from the network.
685        """
686        # Modified b1 rather than conv1 and bn1
687        x = self.b1(x)
688        # Rest of forward remains the same
689        x = F.relu(x)
690        x = self.maxpool(x)
691
692        x = self.layer1(x)
693        x = self.layer2(x)
694        x = self.layer3(x)
695        x = self.layer4(x)
696
697        x = self.avgpool(x)
698        x = torch.flatten(x, 1)
699        x = self.fc(x)
700
701        return x
702
703    def forward(self, x):
704        """Forward pass through the network.
705
706        Parameters
707        ----------
708        x : torch.Tensor
709            Input tensor to the network.
710
711        Returns
712        -------
713        torch.Tensor
714            Output tensor from the network.
715        """
716        return self._forward_impl(x)
717
718
719class ResNetPAIPreFC(nn.Module):
720    """PB-compatible ResNet wrapper.
721
722    All normalization layers should be wrapped in a PAISequential, or other
723    wrapped module. When working with a predefined model the following shows
724    an example of how to create a module for modules_to_replace.
725    """
726
727    def __init__(self, other_resnet):
728        """Initialize ResNetPAI from existing ResNet model.
729
730        Parameters
731        ----------
732        *args : other_resnet : torchvision.models.resnet.ResNet
733            An existing ResNet model to convert to PAI-compatible format.
734        """
735        super(ResNetPAIPreFC, self).__init__()
736
737        # For the most part, just copy the exact values from the original module
738        self._norm_layer = other_resnet._norm_layer
739        self.inplanes = other_resnet.inplanes
740        self.dilation = other_resnet.dilation
741        self.groups = other_resnet.groups
742        self.base_width = other_resnet.base_width
743
744        # For the component to be changed, define a PAISequential with the old
745        # modules included
746        self.conv1 = other_resnet.conv1
747        self.bn1 = other_resnet.bn1
748
749        self.relu = other_resnet.relu
750        self.maxpool = other_resnet.maxpool
751
752        for i in range(1, 5):
753            layer_name = "layer" + str(i)
754            original_layer = getattr(other_resnet, layer_name)
755            pb_layer = self._make_layer_pb(original_layer, other_resnet, i)
756            setattr(self, layer_name, pb_layer)
757
758        self.avgpool = other_resnet.avgpool
759        
760        # Create pre_fc layer with dimensions matching layer4 output (same as fc input)
761        fc_in_features = other_resnet.fc.in_features
762        self.pre_fc = nn.Linear(fc_in_features, fc_in_features)
763        
764        self.fc = other_resnet.fc
765
766    def _make_layer_pb(self, other_block_set, other_resnet, block_id):
767        """Convert ResNet layer blocks to PB-compatible format.
768
769        Parameters
770        ----------
771        other_block_set : torch.vision.models.resnet.any_block
772            A set of blocks from the original ResNet model.
773        other_resnet : torchvision.models.resnet.ResNet
774            The original ResNet model.
775        block_id : int
776            The layer number being converted.
777        Returns
778        -------
779        nn.Sequential
780            A sequential container with the converted blocks.
781        """
782        layers = []
783        for i in range(len(other_block_set)):
784            block_type = type(other_block_set[i])
785            if block_type == resnet_pt.BasicBlock:
786                layers.append(other_block_set[i])
787            elif block_type == resnet_pt.Bottleneck:
788                layers.append(other_block_set[i])
789            else:
790                print(
791                    "Your resnet uses a block type that has not been "
792                    "accounted for. Customization might be required."
793                )
794                layer_name = "layer" + str(block_id)
795                print(type(getattr(other_resnet, layer_name)))
796                pdb.set_trace()
797        return nn.Sequential(*layers)
798
799    def _forward_impl(self, x):
800        """Implementation of the forward pass.
801
802        Parameters
803        ----------
804        x : torch.Tensor
805            Input tensor to the network.
806
807        Returns
808        -------
809        torch.Tensor
810            Output tensor from the network.
811        """
812        # Modified b1 rather than conv1 and bn1
813        x = self.conv1(x)
814        x = self.bn1(x)
815        # Rest of forward remains the same
816        x = F.relu(x)
817        x = self.maxpool(x)
818
819        x = self.layer1(x)
820        x = self.layer2(x)
821        x = self.layer3(x)
822        x = self.layer4(x)
823
824        x = self.avgpool(x)
825        x = torch.flatten(x, 1)
826        x = self.pre_fc(x)
827        x = F.relu(x)
828        x = self.fc(x)
829
830        return x
831
832    def forward(self, x):
833        """Forward pass through the network.
834
835        Parameters
836        ----------
837        x : torch.Tensor
838            Input tensor to the network.
839
840        Returns
841        -------
842        torch.Tensor
843            Output tensor from the network.
844        """
845        return self._forward_impl(x)