perforatedai.network_perforatedai

  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    """Convert a model to use PAI perforated wrappers.
 21
 22    Parameters
 23    ----------
 24    net : nn.Module
 25        Model or submodule to convert.
 26    layer_name : str, optional
 27        Required name when converting a single module directly.
 28
 29    Returns
 30    -------
 31    nn.Module
 32        Converted module tree.
 33    """
 34    # If the net itself has a substitution make that substitution first
 35    if type(net) in GPA.pc.get_modules_to_replace():
 36        net = UPA.replace_predefined_modules(net)
 37    # If the net itself should be converted make the converstion
 38    if type(net) in GPA.pc.get_modules_to_perforate():
 39        if layer_name == "":
 40            print(
 41                "converting a single layer without a name, add a layer_name param to the call"
 42            )
 43            sys.exit(-1)
 44        net = PerforatedModule(net, layer_name)
 45    # Otherwise, check the module recursively if there are other modules to convert
 46    else:
 47        net = UPA.convert_module(net, 0, "", [], [], PerforatedModule, PAITrackedModule)
 48    return net
 49
 50
 51def get_pai_modules(net, depth, seen_ids=None):
 52    """Collect unique PerforatedModule instances from a module tree.
 53
 54    Parameters
 55    ----------
 56    net : nn.Module
 57        Root module to traverse.
 58    depth : int
 59        Current recursion depth.
 60    seen_ids : set or None, optional
 61        Set of module object IDs already collected.
 62
 63    Returns
 64    -------
 65    list
 66        List of unique ``PerforatedModule`` objects.
 67    """
 68    if seen_ids is None:
 69        seen_ids = set()
 70    all_members = net.__dir__()
 71    this_list = []
 72    if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList):
 73        for submodule_id, layer in net.named_children():
 74            if net.get_submodule(submodule_id) is net:
 75                continue
 76            if type(net.get_submodule(submodule_id)) is PerforatedModule:
 77                module = net.get_submodule(submodule_id)
 78                if id(module) in seen_ids:
 79                    continue
 80                seen_ids.add(id(module))
 81                this_list = this_list + [module]
 82            else:
 83                this_list = this_list + get_pai_modules(
 84                    net.get_submodule(submodule_id), depth + 1, seen_ids
 85                )
 86    else:
 87        for member in all_members:
 88            if isinstance(getattr(type(net), member, None), property):
 89                continue
 90            if getattr(net, member, None) is net:
 91                continue
 92            if type(getattr(net, member, None)) is PerforatedModule:
 93                module = getattr(net, member)
 94                if id(module) in seen_ids:
 95                    continue
 96                seen_ids.add(id(module))
 97                this_list = this_list + [module]
 98            elif issubclass(type(getattr(net, member, None)), nn.Module):
 99                this_list = this_list + get_pai_modules(
100                    getattr(net, member), depth + 1, seen_ids
101                )
102    return this_list
103
104
105def load_pai_model_from_dict(net, state_dict):
106    """Load a PAI or plain model from a state dictionary into an unconverted network.
107
108    Parameters
109    ----------
110    net : nn.Module
111        Base network architecture (not yet converted to PerforatedModules).
112    state_dict : dict
113        Serialized model state.
114
115    Returns
116    -------
117    nn.Module
118        Network with loaded state and reconstructed runtime buffers.
119    """
120    # Find which module paths need PAI wrapping from the state dict
121    pai_module_names = set(
122        key[: -len(".num_cycles")]
123        for key in state_dict.keys()
124        if key.endswith(".num_cycles")
125    )
126
127    # Clean up tracked-module scaffolding: for any key containing .main_module.
128    # where the prefix is not a PAI module, strip .main_module. out so the
129    # weights load directly into the plain module. Also drop any module_id keys.
130    cleaned = {}
131    for key, value in state_dict.items():
132        if ".main_module." in key:
133            prefix = key[: key.index(".main_module.")]
134            if prefix not in pai_module_names:
135                key = key.replace(".main_module.", ".", 1)
136        if "module_id" in key.split("."):
137            continue
138        cleaned[key] = value
139    state_dict = cleaned
140
141    if not pai_module_names:
142        net.load_state_dict(state_dict)
143        return net
144
145    # Wrap each identified module as a PerforatedModule in-place
146    for module_name in pai_module_names:
147        parts = module_name.split(".")
148        if len(parts) == 1:
149            parent = net
150            attr = parts[0]
151        else:
152            parent = net.get_submodule(".".join(parts[:-1]))
153            attr = parts[-1]
154        original_module = getattr(parent, attr)
155        setattr(parent, attr, PerforatedModule(original_module, "." + module_name))
156
157    pai_modules = get_pai_modules(net, 0)
158
159    for module in pai_modules:
160        # Set up name to be what will be saved in the state dict
161        module_name = UPA.get_module_base_name(module)
162        # Then instantiate as many Dendrites as were created during training
163        num_cycles = int(state_dict[module_name + ".num_cycles"].item())
164        # extract node index from state_dict
165        nodeCount = 10
166        # also extract view tuple
167        if num_cycles > 0:
168            module.simulate_cycles(num_cycles, nodeCount)
169        if not module.processor is None:
170            processor = copy.deepcopy(module.processor)
171            processor.pre = module.processor.post_n1
172            processor.post = module.processor.post_n2
173            module.processor_array.append(processor)
174        else:
175            module.processor_array.append(None)
176
177        # Create ParameterList for skip_weights based on num_cycles
178        num_params = num_cycles // 2
179        skip_weights_list = nn.ParameterList()
180        for i in range(num_params):
181            param_key = module_name + f".skip_weights.{i}"
182            if param_key in state_dict:
183                param = nn.Parameter(torch.randn(state_dict[param_key].shape))
184                skip_weights_list.append(param)
185        module.skip_weights = skip_weights_list
186
187        module.register_buffer("view_tuple", state_dict[module_name + ".view_tuple"])
188
189    net.load_state_dict(state_dict)
190
191    for module in pai_modules:
192        temp = tuple(module.view_tuple.tolist())
193        del module.view_tuple
194        module.view_tuple = temp
195
196    return net
197
198
199def load_pai_model(net, filename):
200    """Load a saved PAI model file into a network.
201
202    Parameters
203    ----------
204    net : nn.Module
205        Base network architecture.
206    filename : str
207        Path to a safetensors state file.
208
209    Returns
210    -------
211    nn.Module
212        Loaded network.
213    """
214    state_dict = load_file(filename)
215    return load_pai_model_from_dict(net, state_dict)
216
217
218class PerforatedModule(nn.Module):
219    def __init__(self, original_module, name):
220        """Initialize a perforated wrapper around a single module.
221
222        Parameters
223        ----------
224        original_module : nn.Module
225            Original module being wrapped.
226        name : str
227            Qualified module name used for save/load mapping.
228        """
229        super(PerforatedModule, self).__init__()
230        self.name = name
231        self.register_buffer("node_index", torch.tensor(-1))
232        self.register_buffer("num_cycles", torch.tensor(-1))
233        self.register_buffer("view_tuple", torch.tensor(-1))
234        self.processor_array = []
235        self.processor = None
236        self.layer_array = nn.ModuleList([original_module])
237        # If this original module has processing functions save the processor
238        if type(original_module) in GPA.pc.get_modules_with_processing():
239            module_index = GPA.pc.get_modules_with_processing().index(
240                type(original_module)
241            )
242            self.processor = GPA.pc.get_modules_processing_classes()[module_index]()
243        elif (
244            type(original_module).__name__ in GPA.pc.get_module_names_with_processing()
245        ):
246            module_index = GPA.pc.get_module_names_with_processing().index(
247                type(original_module).__name__
248            )
249            self.processor = GPA.pc.get_module_by_name_processing_classes()[
250                module_index
251            ]()
252
253    def simulate_cycles(self, num_cycles, nodeCount):
254        """Expand internal layer/processor lists for stored dendrite cycles.
255
256        Parameters
257        ----------
258        num_cycles : int
259            Number of perforation cycles represented in saved state.
260        nodeCount : int
261            Node count hint kept for interface compatibility.
262
263        Returns
264        -------
265        None
266            This function does not return a value.
267        """
268        for i in range(0, num_cycles, 2):
269            self.layer_array.append(copy.deepcopy(self.layer_array[0]))
270            if not self.processor is None:
271                processor = copy.deepcopy(self.processor)
272                processor.pre = self.processor.pre_d
273                processor.post = self.processor.post_d
274                self.processor_array.append(processor)
275            else:
276                self.processor_array.append(None)
277
278    def process_and_forward(self, *args2, **kwargs2):
279        """Execute one dendrite layer and write output into shared storage.
280
281        Parameters
282        ----------
283        *args2 : tuple
284            Positional values where first entries are layer index and output
285            buffer, followed by layer inputs.
286        **kwargs2 : dict
287            Keyword arguments forwarded to the wrapped layer.
288
289        Returns
290        -------
291        None
292            This function does not return a value.
293        """
294        c = args2[0]
295        dendrite_outs = args2[1]
296        args2 = args2[2:]
297        if self.processor_array[c] != None:
298            out_values = self.processor_array[c].pre(*args2, **kwargs2)
299        out_values = self.layer_array[c](*args2, **kwargs2)
300        if self.processor_array[c] != None:
301            out = self.processor_array[c].post(out_values)
302        else:
303            out = out_values
304        dendrite_outs[c] = out
305
306    def process_and_pre(self, *args, **kwargs):
307        """Run the final pre-dendrite layer pass and cache its output.
308
309        Parameters
310        ----------
311        *args : tuple
312            Positional values with output buffer first, then model inputs.
313        **kwargs : dict
314            Keyword arguments forwarded to the wrapped layer.
315
316        Returns
317        -------
318        None
319            This function does not return a value.
320        """
321        dendrite_outs = args[0]
322        args = args[1:]
323        out = self.layer_array[-1].forward(*args, **kwargs)
324        if not self.processor_array[-1] is None:
325            out = self.processor_array[-1].pre(out)
326        dendrite_outs[len(self.layer_array) - 1] = out
327
328    def forward(self, *args, **kwargs):
329        """Run perforated forward pass with dendrite accumulation.
330
331        Parameters
332        ----------
333        *args : tuple
334            Positional arguments forwarded through wrapped layers.
335        **kwargs : dict
336            Keyword arguments forwarded through wrapped layers.
337
338        Returns
339        -------
340        Any
341            Final model output after optional processor post-processing.
342        """
343        # this is currently false anyway, just remove the doing multi idea
344        doing_multi = doing_threading
345        dendrite_outs = [None] * len(self.layer_array)
346        threads = {}
347        for c in range(0, len(self.layer_array) - 1):
348            args2, kwargs2 = args, kwargs
349            if doing_multi:
350                threads[c] = Thread(
351                    target=self.process_and_forward,
352                    args=(c, dendrite_outs, *args),
353                    kwargs=kwargs,
354                )
355            else:
356                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
357        if doing_multi:
358            threads[len(self.layer_array) - 1] = Thread(
359                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
360            )
361        else:
362            self.process_and_pre(dendrite_outs, *args, **kwargs)
363        if doing_multi:
364            for i in range(len(dendrite_outs)):
365                threads[i].start()
366            for i in range(len(dendrite_outs)):
367                threads[i].join()
368        for out_index in range(0, len(self.layer_array)):
369            current_out = dendrite_outs[out_index]
370
371            if len(self.layer_array) > 1 and hasattr(self, "skip_weights") and len(self.skip_weights) > 0:
372                for in_index in range(0, out_index):
373                    # Use out_index - 1 because skip_weights[0] is never used
374                    current_out = (
375                        current_out
376                        + self.skip_weights[out_index - 1][in_index, :]
377                        .reshape(self.view_tuple)
378                        .to(current_out.device)
379                        * dendrite_outs[in_index]
380                    )
381                if out_index < len(self.layer_array) - 1:
382                    current_out = GPA.pc.get_pai_forward_function()(current_out)
383            dendrite_outs[out_index] = current_out
384        if not self.processor_array[-1] is None:
385            current_out = self.processor_array[-1].post(current_out)
386        return current_out
387
388
389class PAITrackedModule(nn.Module):
390    """Wrapper for modules you don't want to add dendrites to. Ensures all modules are accounted for."""
391
392    def __init__(self, start_module, name):
393        """Initialize PAITrackedModule.
394
395        This function sets up the tracked neuron module to wrap the start_module
396        without adding dendrites.
397
398        Parameters
399        ----------
400        start_module : nn.Module
401            The module to wrap.
402        name : str
403            The name of the neuron module.
404        """
405        super(PAITrackedModule, self).__init__()
406
407        if isinstance(start_module, nn.Module):
408            self.main_module = start_module
409        else:
410            print("start_module must be nn.Module: %s" % name)
411            print(type(start_module))
412            print(start_module)
413            sys.exit(-1)
414        self.name = name
415
416        self.type = "tracked_module"
417
418    def __getattr__(self, name):
419        """Get member variables from the main module.
420
421        Parameters
422        ----------
423        name : str
424            The name of the variable to retrieve.
425        Returns
426        -------
427        The requested variable.
428
429        Notes
430        -----
431        This method first attempts to retrieve the attribute from the PAINeuronModule instance.
432        If it fails, it tries to get the attribute from the wrapped main_module.
433        This allows seamless access to the main module's attributes without modifying original code.
434        """
435        try:
436            return super().__getattr__(name)
437        except AttributeError:
438            return getattr(self.main_module, name)
439
440    def forward(self, *args, **kwargs):
441        """Forward pass for tracked layer.
442
443        Parameters
444        ----------
445        *args : tuple
446            Positional arguments for the forward pass.
447        **kwargs : dict
448            Keyword arguments for the forward pass.
449
450        Returns
451        -------
452        Any
453            The output of the module
454
455        Notes
456        -----
457            The output of this forward function will have the same format as the output
458            of the original module
459        """
460        return self.main_module(*args, **kwargs)
461
462    def __str__(self):
463        """String representation of the layer.
464
465        Parameters
466        ----------
467        None
468
469        Returns
470        -------
471        str
472            String representation of the layer.
473
474        Notes
475        -----
476        Setting for verbose changes level of details in the string output.
477        """
478
479        if GPA.pc.get_verbose():
480            total_string = self.main_module.__str__()
481            total_string = "PAITrackedLayer(" + total_string + ")"
482            return total_string
483        else:
484            total_string = self.main_module.__str__()
485            total_string = "PAITrackedLayer(" + total_string + ")"
486            return total_string
487
488    def __repr__(self):
489        """Representation of the layer."""
490        return self.__str__()