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 PAI-specific module state from a state dictionary.
107
108    Parameters
109    ----------
110    net : nn.Module
111        Converted network containing perforated modules.
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    pai_modules = get_pai_modules(net, 0)
121    if pai_modules == []:
122        print("No PAI modules were found something went wrong with convert network")
123        pdb.set_trace()
124        sys.exit()
125    for module in pai_modules:
126        # Set up name to be what will be saved in the state dict
127        module_name = UPA.get_module_base_name(module)
128        # Then instantiate as many Dendrites as were created during training
129        num_cycles = int(state_dict[module_name + ".num_cycles"].item())
130        # extract node index from state_dict
131        nodeCount = 10
132        # also extract view tuple
133        if num_cycles > 0:
134            module.simulate_cycles(num_cycles, nodeCount)
135        if not module.processor is None:
136            processor = copy.deepcopy(module.processor)
137            processor.pre = module.processor.post_n1
138            processor.post = module.processor.post_n2
139            module.processor_array.append(processor)
140        else:
141            module.processor_array.append(None)
142
143        # Create ParameterList for skip_weights based on num_cycles
144        num_params = num_cycles // 2
145        skip_weights_list = nn.ParameterList()
146        for i in range(num_params):
147            param_key = module_name + f".skip_weights.{i}"
148            if param_key in state_dict:
149                param = nn.Parameter(torch.randn(state_dict[param_key].shape))
150                skip_weights_list.append(param)
151        module.skip_weights = skip_weights_list
152
153        # module.register_buffer('skip_weights', torch.zeros(state_dict[module_name + '.skip_weights'].shape))
154        module.register_buffer("view_tuple", state_dict[module_name + ".view_tuple"])
155
156    net.load_state_dict(state_dict)
157
158    for module in pai_modules:
159        temp = tuple(module.view_tuple.tolist())
160        del module.view_tuple
161        module.view_tuple = temp
162
163    return net
164    # figure out if doing this 'thread' stuff is actually helping at all.
165    # If its not just get rid of it to simplify things.
166    # to test this will have to first get load_pai_model actually set up and working then run a test with and #without threading.
167
168
169def load_pai_model(net, filename):
170    """Load a saved PAI model file into a converted network.
171
172    Parameters
173    ----------
174    net : nn.Module
175        Base network architecture.
176    filename : str
177        Path to a safetensors state file.
178
179    Returns
180    -------
181    nn.Module
182        Loaded PAI network.
183    """
184    net = convert_network(net)
185    state_dict = load_file(filename)
186    return load_pai_model_from_dict(net, state_dict)
187
188
189class PerforatedModule(nn.Module):
190    def __init__(self, original_module, name):
191        """Initialize a perforated wrapper around a single module.
192
193        Parameters
194        ----------
195        original_module : nn.Module
196            Original module being wrapped.
197        name : str
198            Qualified module name used for save/load mapping.
199        """
200        super(PerforatedModule, self).__init__()
201        self.name = name
202        self.register_buffer("node_index", torch.tensor(-1))
203        self.register_buffer("num_cycles", torch.tensor(-1))
204        self.register_buffer("view_tuple", torch.tensor(-1))
205        self.processor_array = []
206        self.processor = None
207        self.layer_array = nn.ModuleList([original_module])
208        # If this original module has processing functions save the processor
209        if type(original_module) in GPA.pc.get_modules_with_processing():
210            module_index = GPA.pc.get_modules_with_processing().index(
211                type(original_module)
212            )
213            self.processor = GPA.pc.get_modules_processing_classes()[module_index]()
214        elif (
215            type(original_module).__name__ in GPA.pc.get_module_names_with_processing()
216        ):
217            module_index = GPA.pc.get_module_names_with_processing().index(
218                type(original_module).__name__
219            )
220            self.processor = GPA.pc.get_module_by_name_processing_classes()[
221                module_index
222            ]()
223
224    def simulate_cycles(self, num_cycles, nodeCount):
225        """Expand internal layer/processor lists for stored dendrite cycles.
226
227        Parameters
228        ----------
229        num_cycles : int
230            Number of perforation cycles represented in saved state.
231        nodeCount : int
232            Node count hint kept for interface compatibility.
233
234        Returns
235        -------
236        None
237            This function does not return a value.
238        """
239        for i in range(0, num_cycles, 2):
240            self.layer_array.append(copy.deepcopy(self.layer_array[0]))
241            if not self.processor is None:
242                processor = copy.deepcopy(self.processor)
243                processor.pre = self.processor.pre_d
244                processor.post = self.processor.post_d
245                self.processor_array.append(processor)
246            else:
247                self.processor_array.append(None)
248
249    def process_and_forward(self, *args2, **kwargs2):
250        """Execute one dendrite layer and write output into shared storage.
251
252        Parameters
253        ----------
254        *args2 : tuple
255            Positional values where first entries are layer index and output
256            buffer, followed by layer inputs.
257        **kwargs2 : dict
258            Keyword arguments forwarded to the wrapped layer.
259
260        Returns
261        -------
262        None
263            This function does not return a value.
264        """
265        c = args2[0]
266        dendrite_outs = args2[1]
267        args2 = args2[2:]
268        if self.processor_array[c] != None:
269            out_values = self.processor_array[c].pre(*args2, **kwargs2)
270        out_values = self.layer_array[c](*args2, **kwargs2)
271        if self.processor_array[c] != None:
272            out = self.processor_array[c].post(out_values)
273        else:
274            out = out_values
275        dendrite_outs[c] = out
276
277    def process_and_pre(self, *args, **kwargs):
278        """Run the final pre-dendrite layer pass and cache its output.
279
280        Parameters
281        ----------
282        *args : tuple
283            Positional values with output buffer first, then model inputs.
284        **kwargs : dict
285            Keyword arguments forwarded to the wrapped layer.
286
287        Returns
288        -------
289        None
290            This function does not return a value.
291        """
292        dendrite_outs = args[0]
293        args = args[1:]
294        out = self.layer_array[-1].forward(*args, **kwargs)
295        if not self.processor_array[-1] is None:
296            out = self.processor_array[-1].pre(out)
297        dendrite_outs[len(self.layer_array) - 1] = out
298
299    def forward(self, *args, **kwargs):
300        """Run perforated forward pass with dendrite accumulation.
301
302        Parameters
303        ----------
304        *args : tuple
305            Positional arguments forwarded through wrapped layers.
306        **kwargs : dict
307            Keyword arguments forwarded through wrapped layers.
308
309        Returns
310        -------
311        Any
312            Final model output after optional processor post-processing.
313        """
314        # this is currently false anyway, just remove the doing multi idea
315        doing_multi = doing_threading
316        dendrite_outs = [None] * len(self.layer_array)
317        threads = {}
318        for c in range(0, len(self.layer_array) - 1):
319            args2, kwargs2 = args, kwargs
320            if doing_multi:
321                threads[c] = Thread(
322                    target=self.process_and_forward,
323                    args=(c, dendrite_outs, *args),
324                    kwargs=kwargs,
325                )
326            else:
327                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
328        if doing_multi:
329            threads[len(self.layer_array) - 1] = Thread(
330                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
331            )
332        else:
333            self.process_and_pre(dendrite_outs, *args, **kwargs)
334        if doing_multi:
335            for i in range(len(dendrite_outs)):
336                threads[i].start()
337            for i in range(len(dendrite_outs)):
338                threads[i].join()
339        for out_index in range(0, len(self.layer_array)):
340            current_out = dendrite_outs[out_index]
341
342            if len(self.layer_array) > 1 and hasattr(self, "skip_weights"):
343                for in_index in range(0, out_index):
344                    # Use out_index - 1 because skip_weights[0] is never used
345                    current_out = (
346                        current_out
347                        + self.skip_weights[out_index - 1][in_index, :]
348                        .reshape(self.view_tuple)
349                        .to(current_out.device)
350                        * dendrite_outs[in_index]
351                    )
352                if out_index < len(self.layer_array) - 1:
353                    current_out = GPA.pc.get_pai_forward_function()(current_out)
354            dendrite_outs[out_index] = current_out
355        if not self.processor_array[-1] is None:
356            current_out = self.processor_array[-1].post(current_out)
357        return current_out
358
359
360class PAITrackedModule(nn.Module):
361    """Wrapper for modules you don't want to add dendrites to. Ensures all modules are accounted for."""
362
363    def __init__(self, start_module, name):
364        """Initialize PAITrackedModule.
365
366        This function sets up the tracked neuron module to wrap the start_module
367        without adding dendrites.
368
369        Parameters
370        ----------
371        start_module : nn.Module
372            The module to wrap.
373        name : str
374            The name of the neuron module.
375        """
376        super(PAITrackedModule, self).__init__()
377
378        if isinstance(start_module, nn.Module):
379            self.main_module = start_module
380        else:
381            print("start_module must be nn.Module: %s" % name)
382            print(type(start_module))
383            print(start_module)
384            sys.exit(-1)
385        self.name = name
386
387        self.type = "tracked_module"
388
389    def __getattr__(self, name):
390        """Get member variables from the main module.
391
392        Parameters
393        ----------
394        name : str
395            The name of the variable to retrieve.
396        Returns
397        -------
398        The requested variable.
399
400        Notes
401        -----
402        This method first attempts to retrieve the attribute from the PAINeuronModule instance.
403        If it fails, it tries to get the attribute from the wrapped main_module.
404        This allows seamless access to the main module's attributes without modifying original code.
405        """
406        try:
407            return super().__getattr__(name)
408        except AttributeError:
409            return getattr(self.main_module, name)
410
411    def forward(self, *args, **kwargs):
412        """Forward pass for tracked layer.
413
414        Parameters
415        ----------
416        *args : tuple
417            Positional arguments for the forward pass.
418        **kwargs : dict
419            Keyword arguments for the forward pass.
420
421        Returns
422        -------
423        Any
424            The output of the module
425
426        Notes
427        -----
428            The output of this forward function will have the same format as the output
429            of the original module
430        """
431        return self.main_module(*args, **kwargs)
432
433    def __str__(self):
434        """String representation of the layer.
435
436        Parameters
437        ----------
438        None
439
440        Returns
441        -------
442        str
443            String representation of the layer.
444
445        Notes
446        -----
447        Setting for verbose changes level of details in the string output.
448        """
449
450        if GPA.pc.get_verbose():
451            total_string = self.main_module.__str__()
452            total_string = "PAITrackedLayer(" + total_string + ")"
453            return total_string
454        else:
455            total_string = self.main_module.__str__()
456            total_string = "PAITrackedLayer(" + total_string + ")"
457            return total_string
458
459    def __repr__(self):
460        """Representation of the layer."""
461        return self.__str__()