perforatedai.network_perforatedai

View Source

  1from perforatedai import globals_perforatedai as GPA
  2from perforatedai import utils_perforatedai as UPA
  3import sys
  4
  5from safetensors.torch import load_file
  6import copy
  7
  8import torch.nn as nn
  9import torch
 10import pdb
 11
 12from threading import Thread
 13
 14
 15doing_threading = False
 16loaded_full_print = False
 17
 18
 19def convert_network(net, layer_name=""):
 20    # If the net itself has a substitution make that substitution first
 21    if type(net) in GPA.pc.get_modules_to_replace():
 22        net = UPA.replace_predefined_modules(net)
 23    # If the net itself should be converted make the converstion
 24    if type(net) in GPA.pc.get_modules_to_perforate():
 25        if layer_name == "":
 26            print(
 27                "converting a single layer without a name, add a layer_name param to the call"
 28            )
 29            sys.exit(-1)
 30        net = PerforatedModule(net, layer_name)
 31    # Otherwise, check the module recursively if there are other modules to convert
 32    else:
 33        net = UPA.convert_module(net, 0, "", [], [], PerforatedModule, PAITrackedModule)
 34    return net
 35
 36
 37def get_pai_modules(net, depth, seen_ids=None):
 38    if seen_ids is None:
 39        seen_ids = set()
 40    all_members = net.__dir__()
 41    this_list = []
 42    if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList):
 43        for submodule_id, layer in net.named_children():
 44            if net.get_submodule(submodule_id) is net:
 45                continue
 46            if type(net.get_submodule(submodule_id)) is PerforatedModule:
 47                module = net.get_submodule(submodule_id)
 48                if id(module) in seen_ids:
 49                    continue
 50                seen_ids.add(id(module))
 51                this_list = this_list + [module]
 52            else:
 53                this_list = this_list + get_pai_modules(
 54                    net.get_submodule(submodule_id), depth + 1, seen_ids
 55                )
 56    else:
 57        for member in all_members:
 58            if isinstance(getattr(type(net), member, None), property):
 59                continue
 60            if getattr(net, member, None) is net:
 61                continue
 62            if type(getattr(net, member, None)) is PerforatedModule:
 63                module = getattr(net, member)
 64                if id(module) in seen_ids:
 65                    continue
 66                seen_ids.add(id(module))
 67                this_list = this_list + [module]
 68            elif issubclass(type(getattr(net, member, None)), nn.Module):
 69                this_list = this_list + get_pai_modules(
 70                    getattr(net, member), depth + 1, seen_ids
 71                )
 72    return this_list
 73
 74
 75def load_pai_model_from_dict(net, state_dict):
 76    pai_modules = get_pai_modules(net, 0)
 77    if pai_modules == []:
 78        print("No PAI modules were found something went wrong with convert network")
 79        pdb.set_trace()
 80        sys.exit()
 81    for module in pai_modules:
 82        # Set up name to be what will be saved in the state dict
 83        module_name = UPA.get_module_base_name(module)
 84        # Then instantiate as many Dendrites as were created during training
 85        num_cycles = int(state_dict[module_name + ".num_cycles"].item())
 86        # extract node index from state_dict
 87        nodeCount = 10
 88        # also extract view tuple
 89        if num_cycles > 0:
 90            module.simulate_cycles(num_cycles, nodeCount)
 91        if not module.processor is None:
 92            processor = copy.deepcopy(module.processor)
 93            processor.pre = module.processor.post_n1
 94            processor.post = module.processor.post_n2
 95            module.processor_array.append(processor)
 96        else:
 97            module.processor_array.append(None)
 98
 99        # Create ParameterList for skip_weights based on num_cycles
100        num_params = num_cycles // 2
101        skip_weights_list = nn.ParameterList()
102        for i in range(num_params):
103            param_key = module_name + f".skip_weights.{i}"
104            if param_key in state_dict:
105                param = nn.Parameter(torch.randn(state_dict[param_key].shape))
106                skip_weights_list.append(param)
107        module.skip_weights = skip_weights_list
108
109        # module.register_buffer('skip_weights', torch.zeros(state_dict[module_name + '.skip_weights'].shape))
110        module.register_buffer("view_tuple", state_dict[module_name + ".view_tuple"])
111
112    net.load_state_dict(state_dict)
113
114    for module in pai_modules:
115        temp = tuple(module.view_tuple.tolist())
116        del module.view_tuple
117        module.view_tuple = temp
118
119    return net
120    # figure out if doing this 'thread' stuff is actually helping at all.
121    # If its not just get rid of it to simplify things.
122    # to test this will have to first get load_pai_model actually set up and working then run a test with and #without threading.
123
124
125def load_pai_model(net, filename):
126    net = convert_network(net)
127    state_dict = load_file(filename)
128    return load_pai_model_from_dict(net, state_dict)
129
130
131class PerforatedModule(nn.Module):
132    def __init__(self, original_module, name):
133        super(PerforatedModule, self).__init__()
134        self.name = name
135        self.register_buffer("node_index", torch.tensor(-1))
136        self.register_buffer("num_cycles", torch.tensor(-1))
137        self.register_buffer("view_tuple", torch.tensor(-1))
138        self.processor_array = []
139        self.processor = None
140        self.layer_array = nn.ModuleList([original_module])
141        # If this original module has processing functions save the processor
142        if type(original_module) in GPA.pc.get_modules_with_processing():
143            module_index = GPA.pc.get_modules_with_processing().index(
144                type(original_module)
145            )
146            self.processor = GPA.pc.get_modules_processing_classes()[module_index]()
147        elif (
148            type(original_module).__name__ in GPA.pc.get_module_names_with_processing()
149        ):
150            module_index = GPA.pc.get_module_names_with_processing().index(
151                type(original_module).__name__
152            )
153            self.processor = GPA.pc.get_module_by_name_processing_classes()[
154                module_index
155            ]()
156
157    def simulate_cycles(self, num_cycles, nodeCount):
158        for i in range(0, num_cycles, 2):
159            self.layer_array.append(copy.deepcopy(self.layer_array[0]))
160            if not self.processor is None:
161                processor = copy.deepcopy(self.processor)
162                processor.pre = self.processor.pre_d
163                processor.post = self.processor.post_d
164                self.processor_array.append(processor)
165            else:
166                self.processor_array.append(None)
167
168    def process_and_forward(self, *args2, **kwargs2):
169        c = args2[0]
170        dendrite_outs = args2[1]
171        args2 = args2[2:]
172        if self.processor_array[c] != None:
173            out_values = self.processor_array[c].pre(*args2, **kwargs2)
174        out_values = self.layer_array[c](*args2, **kwargs2)
175        if self.processor_array[c] != None:
176            out = self.processor_array[c].post(out_values)
177        else:
178            out = out_values
179        dendrite_outs[c] = out
180
181    def process_and_pre(self, *args, **kwargs):
182        dendrite_outs = args[0]
183        args = args[1:]
184        out = self.layer_array[-1].forward(*args, **kwargs)
185        if not self.processor_array[-1] is None:
186            out = self.processor_array[-1].pre(out)
187        dendrite_outs[len(self.layer_array) - 1] = out
188
189    def forward(self, *args, **kwargs):
190        # this is currently false anyway, just remove the doing multi idea
191        doing_multi = doing_threading
192        dendrite_outs = [None] * len(self.layer_array)
193        threads = {}
194        for c in range(0, len(self.layer_array) - 1):
195            args2, kwargs2 = args, kwargs
196            if doing_multi:
197                threads[c] = Thread(
198                    target=self.process_and_forward,
199                    args=(c, dendrite_outs, *args),
200                    kwargs=kwargs,
201                )
202            else:
203                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
204        if doing_multi:
205            threads[len(self.layer_array) - 1] = Thread(
206                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
207            )
208        else:
209            self.process_and_pre(dendrite_outs, *args, **kwargs)
210        if doing_multi:
211            for i in range(len(dendrite_outs)):
212                threads[i].start()
213            for i in range(len(dendrite_outs)):
214                threads[i].join()
215        for out_index in range(0, len(self.layer_array)):
216            current_out = dendrite_outs[out_index]
217
218            if len(self.layer_array) > 1 and hasattr(self, "skip_weights"):
219                for in_index in range(0, out_index):
220                    # Use out_index - 1 because skip_weights[0] is never used
221                    current_out = (
222                        current_out
223                        + self.skip_weights[out_index - 1][in_index, :]
224                        .reshape(self.view_tuple)
225                        .to(current_out.device)
226                        * dendrite_outs[in_index]
227                    )
228                if out_index < len(self.layer_array) - 1:
229                    current_out = GPA.pc.get_pai_forward_function()(current_out)
230            dendrite_outs[out_index] = current_out
231        if not self.processor_array[-1] is None:
232            current_out = self.processor_array[-1].post(current_out)
233        return current_out
234
235
236class PAITrackedModule(nn.Module):
237    """Wrapper for modules you don't want to add dendrites to. Ensures all modules are accounted for."""
238
239    def __init__(self, start_module, name):
240        """Initialize PAITrackedModule.
241
242        This function sets up the tracked neuron module to wrap the start_module
243        without adding dendrites.
244
245        Parameters
246        ----------
247        start_module : nn.Module
248            The module to wrap.
249        name : str
250            The name of the neuron module.
251        """
252        super(PAITrackedModule, self).__init__()
253
254        if isinstance(start_module, nn.Module):
255            self.main_module = start_module
256        else:
257            print("start_module must be nn.Module: %s" % name)
258            print(type(start_module))
259            print(start_module)
260            sys.exit(-1)
261        self.name = name
262
263        self.type = "tracked_module"
264
265    def __getattr__(self, name):
266        """Get member variables from the main module.
267
268        Parameters
269        ----------
270        name : str
271            The name of the variable to retrieve.
272        Returns
273        -------
274        The requested variable.
275
276        Notes
277        -----
278        This method first attempts to retrieve the attribute from the PAINeuronModule instance.
279        If it fails, it tries to get the attribute from the wrapped main_module.
280        This allows seamless access to the main module's attributes without modifying original code.
281        """
282        try:
283            return super().__getattr__(name)
284        except AttributeError:
285            return getattr(self.main_module, name)
286
287    def forward(self, *args, **kwargs):
288        """Forward pass for tracked layer.
289
290        Parameters
291        ----------
292        *args : tuple
293            Positional arguments for the forward pass.
294        **kwargs : dict
295            Keyword arguments for the forward pass.
296
297        Returns
298        -------
299        Any
300            The output of the module
301
302        Notes
303        -----
304            The output of this forward function will have the same format as the output
305            of the original module
306        """
307        return self.main_module(*args, **kwargs)
308
309    def __str__(self):
310        """String representation of the layer.
311
312        Parameters
313        ----------
314        None
315
316        Returns
317        -------
318        str
319            String representation of the layer.
320
321        Notes
322        -----
323        Setting for verbose changes level of details in the string output.
324        """
325
326        if GPA.pc.get_verbose():
327            total_string = self.main_module.__str__()
328            total_string = "PAITrackedLayer(" + total_string + ")"
329            return total_string
330        else:
331            total_string = self.main_module.__str__()
332            total_string = "PAITrackedLayer(" + total_string + ")"
333            return total_string
334
335    def __repr__(self):
336        """Representation of the layer."""
337        return self.__str__()

doing_threading = False

loaded_full_print = False

def convert_network(net, layer_name=''): View Source

20def convert_network(net, layer_name=""):
21    # If the net itself has a substitution make that substitution first
22    if type(net) in GPA.pc.get_modules_to_replace():
23        net = UPA.replace_predefined_modules(net)
24    # If the net itself should be converted make the converstion
25    if type(net) in GPA.pc.get_modules_to_perforate():
26        if layer_name == "":
27            print(
28                "converting a single layer without a name, add a layer_name param to the call"
29            )
30            sys.exit(-1)
31        net = PerforatedModule(net, layer_name)
32    # Otherwise, check the module recursively if there are other modules to convert
33    else:
34        net = UPA.convert_module(net, 0, "", [], [], PerforatedModule, PAITrackedModule)
35    return net

def get_pai_modules(net, depth, seen_ids=None): View Source

38def get_pai_modules(net, depth, seen_ids=None):
39    if seen_ids is None:
40        seen_ids = set()
41    all_members = net.__dir__()
42    this_list = []
43    if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList):
44        for submodule_id, layer in net.named_children():
45            if net.get_submodule(submodule_id) is net:
46                continue
47            if type(net.get_submodule(submodule_id)) is PerforatedModule:
48                module = net.get_submodule(submodule_id)
49                if id(module) in seen_ids:
50                    continue
51                seen_ids.add(id(module))
52                this_list = this_list + [module]
53            else:
54                this_list = this_list + get_pai_modules(
55                    net.get_submodule(submodule_id), depth + 1, seen_ids
56                )
57    else:
58        for member in all_members:
59            if isinstance(getattr(type(net), member, None), property):
60                continue
61            if getattr(net, member, None) is net:
62                continue
63            if type(getattr(net, member, None)) is PerforatedModule:
64                module = getattr(net, member)
65                if id(module) in seen_ids:
66                    continue
67                seen_ids.add(id(module))
68                this_list = this_list + [module]
69            elif issubclass(type(getattr(net, member, None)), nn.Module):
70                this_list = this_list + get_pai_modules(
71                    getattr(net, member), depth + 1, seen_ids
72                )
73    return this_list

def load_pai_model_from_dict(net, state_dict): View Source

 76def load_pai_model_from_dict(net, state_dict):
 77    pai_modules = get_pai_modules(net, 0)
 78    if pai_modules == []:
 79        print("No PAI modules were found something went wrong with convert network")
 80        pdb.set_trace()
 81        sys.exit()
 82    for module in pai_modules:
 83        # Set up name to be what will be saved in the state dict
 84        module_name = UPA.get_module_base_name(module)
 85        # Then instantiate as many Dendrites as were created during training
 86        num_cycles = int(state_dict[module_name + ".num_cycles"].item())
 87        # extract node index from state_dict
 88        nodeCount = 10
 89        # also extract view tuple
 90        if num_cycles > 0:
 91            module.simulate_cycles(num_cycles, nodeCount)
 92        if not module.processor is None:
 93            processor = copy.deepcopy(module.processor)
 94            processor.pre = module.processor.post_n1
 95            processor.post = module.processor.post_n2
 96            module.processor_array.append(processor)
 97        else:
 98            module.processor_array.append(None)
 99
100        # Create ParameterList for skip_weights based on num_cycles
101        num_params = num_cycles // 2
102        skip_weights_list = nn.ParameterList()
103        for i in range(num_params):
104            param_key = module_name + f".skip_weights.{i}"
105            if param_key in state_dict:
106                param = nn.Parameter(torch.randn(state_dict[param_key].shape))
107                skip_weights_list.append(param)
108        module.skip_weights = skip_weights_list
109
110        # module.register_buffer('skip_weights', torch.zeros(state_dict[module_name + '.skip_weights'].shape))
111        module.register_buffer("view_tuple", state_dict[module_name + ".view_tuple"])
112
113    net.load_state_dict(state_dict)
114
115    for module in pai_modules:
116        temp = tuple(module.view_tuple.tolist())
117        del module.view_tuple
118        module.view_tuple = temp
119
120    return net
121    # figure out if doing this 'thread' stuff is actually helping at all.
122    # If its not just get rid of it to simplify things.
123    # to test this will have to first get load_pai_model actually set up and working then run a test with and #without threading.

def load_pai_model(net, filename): View Source

126def load_pai_model(net, filename):
127    net = convert_network(net)
128    state_dict = load_file(filename)
129    return load_pai_model_from_dict(net, state_dict)

class PerforatedModule(torch.nn.modules.module.Module): View Source

132class PerforatedModule(nn.Module):
133    def __init__(self, original_module, name):
134        super(PerforatedModule, self).__init__()
135        self.name = name
136        self.register_buffer("node_index", torch.tensor(-1))
137        self.register_buffer("num_cycles", torch.tensor(-1))
138        self.register_buffer("view_tuple", torch.tensor(-1))
139        self.processor_array = []
140        self.processor = None
141        self.layer_array = nn.ModuleList([original_module])
142        # If this original module has processing functions save the processor
143        if type(original_module) in GPA.pc.get_modules_with_processing():
144            module_index = GPA.pc.get_modules_with_processing().index(
145                type(original_module)
146            )
147            self.processor = GPA.pc.get_modules_processing_classes()[module_index]()
148        elif (
149            type(original_module).__name__ in GPA.pc.get_module_names_with_processing()
150        ):
151            module_index = GPA.pc.get_module_names_with_processing().index(
152                type(original_module).__name__
153            )
154            self.processor = GPA.pc.get_module_by_name_processing_classes()[
155                module_index
156            ]()
157
158    def simulate_cycles(self, num_cycles, nodeCount):
159        for i in range(0, num_cycles, 2):
160            self.layer_array.append(copy.deepcopy(self.layer_array[0]))
161            if not self.processor is None:
162                processor = copy.deepcopy(self.processor)
163                processor.pre = self.processor.pre_d
164                processor.post = self.processor.post_d
165                self.processor_array.append(processor)
166            else:
167                self.processor_array.append(None)
168
169    def process_and_forward(self, *args2, **kwargs2):
170        c = args2[0]
171        dendrite_outs = args2[1]
172        args2 = args2[2:]
173        if self.processor_array[c] != None:
174            out_values = self.processor_array[c].pre(*args2, **kwargs2)
175        out_values = self.layer_array[c](*args2, **kwargs2)
176        if self.processor_array[c] != None:
177            out = self.processor_array[c].post(out_values)
178        else:
179            out = out_values
180        dendrite_outs[c] = out
181
182    def process_and_pre(self, *args, **kwargs):
183        dendrite_outs = args[0]
184        args = args[1:]
185        out = self.layer_array[-1].forward(*args, **kwargs)
186        if not self.processor_array[-1] is None:
187            out = self.processor_array[-1].pre(out)
188        dendrite_outs[len(self.layer_array) - 1] = out
189
190    def forward(self, *args, **kwargs):
191        # this is currently false anyway, just remove the doing multi idea
192        doing_multi = doing_threading
193        dendrite_outs = [None] * len(self.layer_array)
194        threads = {}
195        for c in range(0, len(self.layer_array) - 1):
196            args2, kwargs2 = args, kwargs
197            if doing_multi:
198                threads[c] = Thread(
199                    target=self.process_and_forward,
200                    args=(c, dendrite_outs, *args),
201                    kwargs=kwargs,
202                )
203            else:
204                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
205        if doing_multi:
206            threads[len(self.layer_array) - 1] = Thread(
207                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
208            )
209        else:
210            self.process_and_pre(dendrite_outs, *args, **kwargs)
211        if doing_multi:
212            for i in range(len(dendrite_outs)):
213                threads[i].start()
214            for i in range(len(dendrite_outs)):
215                threads[i].join()
216        for out_index in range(0, len(self.layer_array)):
217            current_out = dendrite_outs[out_index]
218
219            if len(self.layer_array) > 1 and hasattr(self, "skip_weights"):
220                for in_index in range(0, out_index):
221                    # Use out_index - 1 because skip_weights[0] is never used
222                    current_out = (
223                        current_out
224                        + self.skip_weights[out_index - 1][in_index, :]
225                        .reshape(self.view_tuple)
226                        .to(current_out.device)
227                        * dendrite_outs[in_index]
228                    )
229                if out_index < len(self.layer_array) - 1:
230                    current_out = GPA.pc.get_pai_forward_function()(current_out)
231            dendrite_outs[out_index] = current_out
232        if not self.processor_array[-1] is None:
233            current_out = self.processor_array[-1].post(current_out)
234        return current_out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

PerforatedModule(original_module, name) View Source

133    def __init__(self, original_module, name):
134        super(PerforatedModule, self).__init__()
135        self.name = name
136        self.register_buffer("node_index", torch.tensor(-1))
137        self.register_buffer("num_cycles", torch.tensor(-1))
138        self.register_buffer("view_tuple", torch.tensor(-1))
139        self.processor_array = []
140        self.processor = None
141        self.layer_array = nn.ModuleList([original_module])
142        # If this original module has processing functions save the processor
143        if type(original_module) in GPA.pc.get_modules_with_processing():
144            module_index = GPA.pc.get_modules_with_processing().index(
145                type(original_module)
146            )
147            self.processor = GPA.pc.get_modules_processing_classes()[module_index]()
148        elif (
149            type(original_module).__name__ in GPA.pc.get_module_names_with_processing()
150        ):
151            module_index = GPA.pc.get_module_names_with_processing().index(
152                type(original_module).__name__
153            )
154            self.processor = GPA.pc.get_module_by_name_processing_classes()[
155                module_index
156            ]()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

name

processor_array

processor

layer_array

def simulate_cycles(self, num_cycles, nodeCount): View Source

158    def simulate_cycles(self, num_cycles, nodeCount):
159        for i in range(0, num_cycles, 2):
160            self.layer_array.append(copy.deepcopy(self.layer_array[0]))
161            if not self.processor is None:
162                processor = copy.deepcopy(self.processor)
163                processor.pre = self.processor.pre_d
164                processor.post = self.processor.post_d
165                self.processor_array.append(processor)
166            else:
167                self.processor_array.append(None)

def process_and_forward(self, *args2, **kwargs2): View Source

169    def process_and_forward(self, *args2, **kwargs2):
170        c = args2[0]
171        dendrite_outs = args2[1]
172        args2 = args2[2:]
173        if self.processor_array[c] != None:
174            out_values = self.processor_array[c].pre(*args2, **kwargs2)
175        out_values = self.layer_array[c](*args2, **kwargs2)
176        if self.processor_array[c] != None:
177            out = self.processor_array[c].post(out_values)
178        else:
179            out = out_values
180        dendrite_outs[c] = out

def process_and_pre(self, *args, **kwargs): View Source

182    def process_and_pre(self, *args, **kwargs):
183        dendrite_outs = args[0]
184        args = args[1:]
185        out = self.layer_array[-1].forward(*args, **kwargs)
186        if not self.processor_array[-1] is None:
187            out = self.processor_array[-1].pre(out)
188        dendrite_outs[len(self.layer_array) - 1] = out

def forward(self, *args, **kwargs): View Source

190    def forward(self, *args, **kwargs):
191        # this is currently false anyway, just remove the doing multi idea
192        doing_multi = doing_threading
193        dendrite_outs = [None] * len(self.layer_array)
194        threads = {}
195        for c in range(0, len(self.layer_array) - 1):
196            args2, kwargs2 = args, kwargs
197            if doing_multi:
198                threads[c] = Thread(
199                    target=self.process_and_forward,
200                    args=(c, dendrite_outs, *args),
201                    kwargs=kwargs,
202                )
203            else:
204                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
205        if doing_multi:
206            threads[len(self.layer_array) - 1] = Thread(
207                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
208            )
209        else:
210            self.process_and_pre(dendrite_outs, *args, **kwargs)
211        if doing_multi:
212            for i in range(len(dendrite_outs)):
213                threads[i].start()
214            for i in range(len(dendrite_outs)):
215                threads[i].join()
216        for out_index in range(0, len(self.layer_array)):
217            current_out = dendrite_outs[out_index]
218
219            if len(self.layer_array) > 1 and hasattr(self, "skip_weights"):
220                for in_index in range(0, out_index):
221                    # Use out_index - 1 because skip_weights[0] is never used
222                    current_out = (
223                        current_out
224                        + self.skip_weights[out_index - 1][in_index, :]
225                        .reshape(self.view_tuple)
226                        .to(current_out.device)
227                        * dendrite_outs[in_index]
228                    )
229                if out_index < len(self.layer_array) - 1:
230                    current_out = GPA.pc.get_pai_forward_function()(current_out)
231            dendrite_outs[out_index] = current_out
232        if not self.processor_array[-1] is None:
233            current_out = self.processor_array[-1].post(current_out)
234        return current_out

Define the computation performed at every call.

Should be overridden by all subclasses.

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class PAITrackedModule(torch.nn.modules.module.Module): View Source

237class PAITrackedModule(nn.Module):
238    """Wrapper for modules you don't want to add dendrites to. Ensures all modules are accounted for."""
239
240    def __init__(self, start_module, name):
241        """Initialize PAITrackedModule.
242
243        This function sets up the tracked neuron module to wrap the start_module
244        without adding dendrites.
245
246        Parameters
247        ----------
248        start_module : nn.Module
249            The module to wrap.
250        name : str
251            The name of the neuron module.
252        """
253        super(PAITrackedModule, self).__init__()
254
255        if isinstance(start_module, nn.Module):
256            self.main_module = start_module
257        else:
258            print("start_module must be nn.Module: %s" % name)
259            print(type(start_module))
260            print(start_module)
261            sys.exit(-1)
262        self.name = name
263
264        self.type = "tracked_module"
265
266    def __getattr__(self, name):
267        """Get member variables from the main module.
268
269        Parameters
270        ----------
271        name : str
272            The name of the variable to retrieve.
273        Returns
274        -------
275        The requested variable.
276
277        Notes
278        -----
279        This method first attempts to retrieve the attribute from the PAINeuronModule instance.
280        If it fails, it tries to get the attribute from the wrapped main_module.
281        This allows seamless access to the main module's attributes without modifying original code.
282        """
283        try:
284            return super().__getattr__(name)
285        except AttributeError:
286            return getattr(self.main_module, name)
287
288    def forward(self, *args, **kwargs):
289        """Forward pass for tracked layer.
290
291        Parameters
292        ----------
293        *args : tuple
294            Positional arguments for the forward pass.
295        **kwargs : dict
296            Keyword arguments for the forward pass.
297
298        Returns
299        -------
300        Any
301            The output of the module
302
303        Notes
304        -----
305            The output of this forward function will have the same format as the output
306            of the original module
307        """
308        return self.main_module(*args, **kwargs)
309
310    def __str__(self):
311        """String representation of the layer.
312
313        Parameters
314        ----------
315        None
316
317        Returns
318        -------
319        str
320            String representation of the layer.
321
322        Notes
323        -----
324        Setting for verbose changes level of details in the string output.
325        """
326
327        if GPA.pc.get_verbose():
328            total_string = self.main_module.__str__()
329            total_string = "PAITrackedLayer(" + total_string + ")"
330            return total_string
331        else:
332            total_string = self.main_module.__str__()
333            total_string = "PAITrackedLayer(" + total_string + ")"
334            return total_string
335
336    def __repr__(self):
337        """Representation of the layer."""
338        return self.__str__()

Wrapper for modules you don't want to add dendrites to. Ensures all modules are accounted for.

PAITrackedModule(start_module, name) View Source

240    def __init__(self, start_module, name):
241        """Initialize PAITrackedModule.
242
243        This function sets up the tracked neuron module to wrap the start_module
244        without adding dendrites.
245
246        Parameters
247        ----------
248        start_module : nn.Module
249            The module to wrap.
250        name : str
251            The name of the neuron module.
252        """
253        super(PAITrackedModule, self).__init__()
254
255        if isinstance(start_module, nn.Module):
256            self.main_module = start_module
257        else:
258            print("start_module must be nn.Module: %s" % name)
259            print(type(start_module))
260            print(start_module)
261            sys.exit(-1)
262        self.name = name
263
264        self.type = "tracked_module"

Initialize PAITrackedModule.

This function sets up the tracked neuron module to wrap the start_module without adding dendrites.

Parameters

start_module (nn.Module): The module to wrap.
name (str): The name of the neuron module.

name

def type(self, dst_type: torch.dtype | str) -> Self: View Source

1167    def type(self, dst_type: dtype | str) -> Self:
1168        r"""Casts all parameters and buffers to :attr:`dst_type`.
1169
1170        .. note::
1171            This method modifies the module in-place.
1172
1173        Args:
1174            dst_type (type or string): the desired type
1175
1176        Returns:
1177            Module: self
1178        """
1179        return self._apply(lambda t: t.type(dst_type))

Casts all parameters and buffers to dst_type.

This method modifies the module in-place.

Args: dst_type (type or string): the desired type

Returns: Module: self

def forward(self, *args, **kwargs): View Source

288    def forward(self, *args, **kwargs):
289        """Forward pass for tracked layer.
290
291        Parameters
292        ----------
293        *args : tuple
294            Positional arguments for the forward pass.
295        **kwargs : dict
296            Keyword arguments for the forward pass.
297
298        Returns
299        -------
300        Any
301            The output of the module
302
303        Notes
304        -----
305            The output of this forward function will have the same format as the output
306            of the original module
307        """
308        return self.main_module(*args, **kwargs)

Forward pass for tracked layer.

Parameters

*args (tuple): Positional arguments for the forward pass.
**kwargs (dict): Keyword arguments for the forward pass.

Returns

Any: The output of the module

Notes

The output of this forward function will have the same format as the output of the original module