perforatedai.clean_perforatedai 

  1# Copyright (c) 2025 Perforated AI
  2from perforatedai import globals_perforatedai as GPA
  3
  4import copy
  5
  6import torch.nn as nn
  7import torch
  8import pdb
  9
 10from threading import Thread
 11
 12doing_threading = False
 13
 14# This is one implimentation of the forward function of PAI modules that 
 15# has an option to use python threading
 16class PAIModulePyThread(nn.Module):
 17    def __init__(self, original_module):
 18        super(PAIModulePyThread, self).__init__()
 19        self.layer_array = original_module.layer_array
 20        self.processor_array = original_module.processor_array
 21        # Remove the unused first index (skip_weights[0] is never used)
 22        if hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) > 1:
 23            self.skip_weights = original_module.skip_weights[1:]
 24        elif hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) == 1:
 25            # Only one element, don't create skip_weights
 26            pass
 27        self.register_buffer("node_index", original_module.node_index.clone().detach())
 28        self.register_buffer("num_cycles", original_module.num_cycles)
 29        self.register_buffer("view_tuple", original_module.view_tuple)
 30
 31    def process_and_forward(self, *args2, **kwargs2):
 32        c = args2[0]
 33        dendrite_outs = args2[1]
 34        args2 = args2[2:]
 35        if self.processor_array[c] != None:
 36            args2, kwargs2 = self.processor_array[c].pre(*args2, **kwargs2)
 37        out_values = self.layer_array[c](*args2, **kwargs2)
 38        if self.processor_array[c] != None:
 39            out = self.processor_array[c].post(out_values)
 40        else:
 41            out = out_values
 42        dendrite_outs[c] = out
 43
 44    def process_and_pre(self, *args, **kwargs):
 45        dendrite_outs = args[0]
 46        args = args[1:]
 47        out = self.layer_array[-1].forward(*args, **kwargs)
 48        if not self.processor_array[-1] is None:
 49            out = self.processor_array[-1].pre(out)
 50        dendrite_outs[len(self.layer_array) - 1] = out
 51
 52    def forward(self, *args, **kwargs):
 53        # this is currently false anyway, just remove the doing multi idea
 54        doing_multi = doing_threading
 55        dendrite_outs = [None] * len(self.layer_array)
 56        threads = {}
 57        for c in range(0, len(self.layer_array) - 1):
 58            args2, kwargs2 = args, kwargs
 59            if doing_multi:
 60                threads[c] = Thread(
 61                    target=self.process_and_forward,
 62                    args=(c, dendrite_outs, *args),
 63                    kwargs=kwargs,
 64                )
 65            else:
 66                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
 67        if doing_multi:
 68            threads[len(self.layer_array) - 1] = Thread(
 69                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
 70            )
 71        else:
 72            self.process_and_pre(dendrite_outs, *args, **kwargs)
 73        if doing_multi:
 74            for i in range(len(dendrite_outs)):
 75                threads[i].start()
 76            for i in range(len(dendrite_outs)):
 77                threads[i].join()
 78        for out_index in range(0, len(self.layer_array)):
 79            current_out = dendrite_outs[out_index]
 80            if len(self.layer_array) > 1 and hasattr(self, 'skip_weights'):
 81                for in_index in range(0, out_index):
 82                    # Use out_index - 1 because skip_weights[0] was removed
 83                    skip_weight = self.skip_weights[out_index - 1][in_index, :]
 84                    # Use cached Python tuple instead of .tolist() during forward
 85                    skip_weight = skip_weight.view(self.view_tuple.tolist())
 86                    current_out = current_out + (
 87                        skip_weight.to(current_out.device)
 88                        * dendrite_outs[in_index]
 89                    )
 90                if out_index < len(self.layer_array) - 1:
 91                    current_out = GPA.pc.get_pai_forward_function()(current_out)
 92            dendrite_outs[out_index] = current_out
 93        if not self.processor_array[-1] is None:
 94            current_out = self.processor_array[-1].post(current_out)
 95        return current_out
 96
 97
 98def get_pretrained_pai_attr(pretrained_dendrite, member):
 99    if pretrained_dendrite is None:
100        return None
101    else:
102        return getattr(pretrained_dendrite, member)
103
104
105def get_pretrained_pai_var(pretrained_dendrite, submodule_id):
106    if pretrained_dendrite is None:
107        return None
108    else:
109        return pretrained_dendrite.get_submodule(submodule_id)
110
111ModuleType = PAIModulePyThread
112doing_threading = False
113
114def make_module(module):
115    return ModuleType(module)
116
117# This Refreshes a PAI network with the PyThread Module
118def refresh_pai(net, depth, name_so_far, converted_list):
119    if GPA.pc.get_extra_verbose():
120        print("CL calling convert on %s depth %d" % (net, depth))
121        print(
122            "CL calling convert on %s: %s, depth %d"
123            % (name_so_far, type(net).__name__, depth)
124        )
125    if type(net) is ModuleType:
126        if GPA.pc.get_extra_verbose():
127            print(
128                "this is only being called because something in your model is pointed to twice by two different variables.  Highest thing on the list is one of the duplicates"
129            )
130        return net
131    all_members = net.__dir__()
132    if (
133        issubclass(type(net), nn.Sequential)
134        or issubclass(type(net), nn.ModuleList)
135        or issubclass(type(net), list)
136    ):
137        for submodule_id, layer in net.named_children():
138            if net != net.get_submodule(submodule_id):
139                converted_list += [name_so_far + "[" + str(submodule_id) + "]"]
140                setattr(
141                    net,
142                    submodule_id,
143                    refresh_pai(
144                        net.get_submodule(submodule_id),
145                        depth + 1,
146                        name_so_far + "[" + str(submodule_id) + "]",
147                        converted_list,
148                    ),
149                )
150            if type(net.get_submodule(submodule_id)).__name__ == "PAILayer":
151                setattr(
152                    net,
153                    submodule_id,
154                    make_module(get_pretrained_pai_var(net, submodule_id)),
155                )
156    elif type(net) in GPA.pc.get_modules_to_track():
157        return net
158    else:
159        for member in all_members:
160            if isinstance(getattr(type(net), member, None), property):
161                continue
162            try:
163                getattr(net, member, None)
164            except:
165                continue
166            sub_name = name_so_far + "." + member
167
168            if member == "device" or member == "dtype":
169                continue
170            if sub_name in GPA.pc.get_module_names_to_not_save():
171                continue
172            if name_so_far == "":
173                if (
174                    sub_name in GPA.pc.get_module_names_to_not_save()
175                    or sub_name in converted_list
176                ):
177                    if GPA.pc.get_extra_verbose():
178                        print("Skipping %s during save" % sub_name)
179                    continue
180
181            if (
182                issubclass(type(getattr(net, member, None)), nn.Module)
183                or member == "layer_array"
184            ):
185                converted_list += [sub_name]
186                if net != getattr(net, member):
187                    setattr(
188                        net,
189                        member,
190                        refresh_pai(
191                            getattr(net, member), depth + 1, sub_name, converted_list
192                        ),
193                    )
194            if type(getattr(net, member, None)).__name__ == "PAILayer":
195                setattr(net, member, make_module(get_pretrained_pai_attr(net, member)))
196    if type(net).__name__ == "PAILayer":
197        net = make_module(net)
198    return net
199
200def refresh_net(pretrained_dendrite):
201
202    net = refresh_pai(pretrained_dendrite, 0, "", [])
203    return net
doing_threading = False
class PAIModulePyThread(torch.nn.modules.module.Module): 
17class PAIModulePyThread(nn.Module):
18    def __init__(self, original_module):
19        super(PAIModulePyThread, self).__init__()
20        self.layer_array = original_module.layer_array
21        self.processor_array = original_module.processor_array
22        # Remove the unused first index (skip_weights[0] is never used)
23        if hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) > 1:
24            self.skip_weights = original_module.skip_weights[1:]
25        elif hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) == 1:
26            # Only one element, don't create skip_weights
27            pass
28        self.register_buffer("node_index", original_module.node_index.clone().detach())
29        self.register_buffer("num_cycles", original_module.num_cycles)
30        self.register_buffer("view_tuple", original_module.view_tuple)
31
32    def process_and_forward(self, *args2, **kwargs2):
33        c = args2[0]
34        dendrite_outs = args2[1]
35        args2 = args2[2:]
36        if self.processor_array[c] != None:
37            args2, kwargs2 = self.processor_array[c].pre(*args2, **kwargs2)
38        out_values = self.layer_array[c](*args2, **kwargs2)
39        if self.processor_array[c] != None:
40            out = self.processor_array[c].post(out_values)
41        else:
42            out = out_values
43        dendrite_outs[c] = out
44
45    def process_and_pre(self, *args, **kwargs):
46        dendrite_outs = args[0]
47        args = args[1:]
48        out = self.layer_array[-1].forward(*args, **kwargs)
49        if not self.processor_array[-1] is None:
50            out = self.processor_array[-1].pre(out)
51        dendrite_outs[len(self.layer_array) - 1] = out
52
53    def forward(self, *args, **kwargs):
54        # this is currently false anyway, just remove the doing multi idea
55        doing_multi = doing_threading
56        dendrite_outs = [None] * len(self.layer_array)
57        threads = {}
58        for c in range(0, len(self.layer_array) - 1):
59            args2, kwargs2 = args, kwargs
60            if doing_multi:
61                threads[c] = Thread(
62                    target=self.process_and_forward,
63                    args=(c, dendrite_outs, *args),
64                    kwargs=kwargs,
65                )
66            else:
67                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
68        if doing_multi:
69            threads[len(self.layer_array) - 1] = Thread(
70                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
71            )
72        else:
73            self.process_and_pre(dendrite_outs, *args, **kwargs)
74        if doing_multi:
75            for i in range(len(dendrite_outs)):
76                threads[i].start()
77            for i in range(len(dendrite_outs)):
78                threads[i].join()
79        for out_index in range(0, len(self.layer_array)):
80            current_out = dendrite_outs[out_index]
81            if len(self.layer_array) > 1 and hasattr(self, 'skip_weights'):
82                for in_index in range(0, out_index):
83                    # Use out_index - 1 because skip_weights[0] was removed
84                    skip_weight = self.skip_weights[out_index - 1][in_index, :]
85                    # Use cached Python tuple instead of .tolist() during forward
86                    skip_weight = skip_weight.view(self.view_tuple.tolist())
87                    current_out = current_out + (
88                        skip_weight.to(current_out.device)
89                        * dendrite_outs[in_index]
90                    )
91                if out_index < len(self.layer_array) - 1:
92                    current_out = GPA.pc.get_pai_forward_function()(current_out)
93            dendrite_outs[out_index] = current_out
94        if not self.processor_array[-1] is None:
95            current_out = self.processor_array[-1].post(current_out)
96        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

PAIModulePyThread(original_module) 
18    def __init__(self, original_module):
19        super(PAIModulePyThread, self).__init__()
20        self.layer_array = original_module.layer_array
21        self.processor_array = original_module.processor_array
22        # Remove the unused first index (skip_weights[0] is never used)
23        if hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) > 1:
24            self.skip_weights = original_module.skip_weights[1:]
25        elif hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) == 1:
26            # Only one element, don't create skip_weights
27            pass
28        self.register_buffer("node_index", original_module.node_index.clone().detach())
29        self.register_buffer("num_cycles", original_module.num_cycles)
30        self.register_buffer("view_tuple", original_module.view_tuple)

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

layer_array
processor_array
def process_and_forward(self, *args2, **kwargs2): 
32    def process_and_forward(self, *args2, **kwargs2):
33        c = args2[0]
34        dendrite_outs = args2[1]
35        args2 = args2[2:]
36        if self.processor_array[c] != None:
37            args2, kwargs2 = self.processor_array[c].pre(*args2, **kwargs2)
38        out_values = self.layer_array[c](*args2, **kwargs2)
39        if self.processor_array[c] != None:
40            out = self.processor_array[c].post(out_values)
41        else:
42            out = out_values
43        dendrite_outs[c] = out
def process_and_pre(self, *args, **kwargs): 
45    def process_and_pre(self, *args, **kwargs):
46        dendrite_outs = args[0]
47        args = args[1:]
48        out = self.layer_array[-1].forward(*args, **kwargs)
49        if not self.processor_array[-1] is None:
50            out = self.processor_array[-1].pre(out)
51        dendrite_outs[len(self.layer_array) - 1] = out
def forward(self, *args, **kwargs): 
53    def forward(self, *args, **kwargs):
54        # this is currently false anyway, just remove the doing multi idea
55        doing_multi = doing_threading
56        dendrite_outs = [None] * len(self.layer_array)
57        threads = {}
58        for c in range(0, len(self.layer_array) - 1):
59            args2, kwargs2 = args, kwargs
60            if doing_multi:
61                threads[c] = Thread(
62                    target=self.process_and_forward,
63                    args=(c, dendrite_outs, *args),
64                    kwargs=kwargs,
65                )
66            else:
67                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
68        if doing_multi:
69            threads[len(self.layer_array) - 1] = Thread(
70                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
71            )
72        else:
73            self.process_and_pre(dendrite_outs, *args, **kwargs)
74        if doing_multi:
75            for i in range(len(dendrite_outs)):
76                threads[i].start()
77            for i in range(len(dendrite_outs)):
78                threads[i].join()
79        for out_index in range(0, len(self.layer_array)):
80            current_out = dendrite_outs[out_index]
81            if len(self.layer_array) > 1 and hasattr(self, 'skip_weights'):
82                for in_index in range(0, out_index):
83                    # Use out_index - 1 because skip_weights[0] was removed
84                    skip_weight = self.skip_weights[out_index - 1][in_index, :]
85                    # Use cached Python tuple instead of .tolist() during forward
86                    skip_weight = skip_weight.view(self.view_tuple.tolist())
87                    current_out = current_out + (
88                        skip_weight.to(current_out.device)
89                        * dendrite_outs[in_index]
90                    )
91                if out_index < len(self.layer_array) - 1:
92                    current_out = GPA.pc.get_pai_forward_function()(current_out)
93            dendrite_outs[out_index] = current_out
94        if not self.processor_array[-1] is None:
95            current_out = self.processor_array[-1].post(current_out)
96        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.

def get_pretrained_pai_attr(pretrained_dendrite, member): 
 99def get_pretrained_pai_attr(pretrained_dendrite, member):
100    if pretrained_dendrite is None:
101        return None
102    else:
103        return getattr(pretrained_dendrite, member)
def get_pretrained_pai_var(pretrained_dendrite, submodule_id): 
106def get_pretrained_pai_var(pretrained_dendrite, submodule_id):
107    if pretrained_dendrite is None:
108        return None
109    else:
110        return pretrained_dendrite.get_submodule(submodule_id)
ModuleType = <class 'PAIModulePyThread'>
def make_module(module): 
115def make_module(module):
116    return ModuleType(module)
def refresh_pai(net, depth, name_so_far, converted_list): 
119def refresh_pai(net, depth, name_so_far, converted_list):
120    if GPA.pc.get_extra_verbose():
121        print("CL calling convert on %s depth %d" % (net, depth))
122        print(
123            "CL calling convert on %s: %s, depth %d"
124            % (name_so_far, type(net).__name__, depth)
125        )
126    if type(net) is ModuleType:
127        if GPA.pc.get_extra_verbose():
128            print(
129                "this is only being called because something in your model is pointed to twice by two different variables.  Highest thing on the list is one of the duplicates"
130            )
131        return net
132    all_members = net.__dir__()
133    if (
134        issubclass(type(net), nn.Sequential)
135        or issubclass(type(net), nn.ModuleList)
136        or issubclass(type(net), list)
137    ):
138        for submodule_id, layer in net.named_children():
139            if net != net.get_submodule(submodule_id):
140                converted_list += [name_so_far + "[" + str(submodule_id) + "]"]
141                setattr(
142                    net,
143                    submodule_id,
144                    refresh_pai(
145                        net.get_submodule(submodule_id),
146                        depth + 1,
147                        name_so_far + "[" + str(submodule_id) + "]",
148                        converted_list,
149                    ),
150                )
151            if type(net.get_submodule(submodule_id)).__name__ == "PAILayer":
152                setattr(
153                    net,
154                    submodule_id,
155                    make_module(get_pretrained_pai_var(net, submodule_id)),
156                )
157    elif type(net) in GPA.pc.get_modules_to_track():
158        return net
159    else:
160        for member in all_members:
161            if isinstance(getattr(type(net), member, None), property):
162                continue
163            try:
164                getattr(net, member, None)
165            except:
166                continue
167            sub_name = name_so_far + "." + member
168
169            if member == "device" or member == "dtype":
170                continue
171            if sub_name in GPA.pc.get_module_names_to_not_save():
172                continue
173            if name_so_far == "":
174                if (
175                    sub_name in GPA.pc.get_module_names_to_not_save()
176                    or sub_name in converted_list
177                ):
178                    if GPA.pc.get_extra_verbose():
179                        print("Skipping %s during save" % sub_name)
180                    continue
181
182            if (
183                issubclass(type(getattr(net, member, None)), nn.Module)
184                or member == "layer_array"
185            ):
186                converted_list += [sub_name]
187                if net != getattr(net, member):
188                    setattr(
189                        net,
190                        member,
191                        refresh_pai(
192                            getattr(net, member), depth + 1, sub_name, converted_list
193                        ),
194                    )
195            if type(getattr(net, member, None)).__name__ == "PAILayer":
196                setattr(net, member, make_module(get_pretrained_pai_attr(net, member)))
197    if type(net).__name__ == "PAILayer":
198        net = make_module(net)
199    return net
def refresh_net(pretrained_dendrite): 
201def refresh_net(pretrained_dendrite):
202
203    net = refresh_pai(pretrained_dendrite, 0, "", [])
204    return net