perforatedai.blockwise_perforatedai 

  1# Copyright (c) 2025 Perforated AI
  2
  3from perforatedai import globals_perforatedai as GPA
  4from perforatedai import modules_perforatedai as PA
  5from perforatedai import utils_perforatedai as UPA
  6import torch.nn as nn
  7import torch
  8import pdb
  9import numpy as np
 10import string
 11import copy
 12
 13# This is the cleaner inference version of PAI modules
 14class PAILayer(nn.Module):
 15    def __init__(
 16        self,
 17        layer_array,
 18        processor_array,
 19        dendrites_to_top,
 20        dendrites_to_dendrites,
 21        node_index,
 22        num_cycles,
 23        view_tuple,
 24    ):
 25        super(PAILayer, self).__init__()
 26        self.layer_array = layer_array
 27        self.register_buffer("num_cycles", num_cycles)
 28        self.register_buffer("view_tuple", torch.tensor(view_tuple))
 29
 30        self.processor_array = processor_array
 31        if dendrites_to_dendrites:
 32            self.skip_weights = dendrites_to_dendrites
 33        else:
 34            """
 35            This will only be the case if there is less than 2 dendrites, in these cases an empty array
 36            should still be added so that dendrites_to_top is included at the correct index
 37            """
 38            self.skip_weights = nn.ParameterList()
 39        if dendrites_to_top:
 40            self.skip_weights.append(dendrites_to_top[len(dendrites_to_top) - 1])
 41        else:
 42            self.skip_weights = nn.ParameterList()
 43        
 44        # Delete skip_weights if it's empty (only 1 layer, no skip connections)
 45        if len(self.skip_weights) == 0:
 46            delattr(self, 'skip_weights')
 47
 48        self.node_index = node_index
 49        self.internal_nonlinearity = GPA.pc.get_pai_forward_function()
 50
 51def unWrap_params(model):
 52    for p in model.parameters():
 53        if "wrapped" in p.__dir__():
 54            del p.wrapped
 55
 56# This converts one training PAI module into an inference PAI module
 57def convert_to_pai_layer_block(pretrained_dendrite):
 58    unWrap_params(pretrained_dendrite)
 59    layer_array = []
 60    processor_array = []
 61    for layer_id in range(len(pretrained_dendrite.dendrite_module.layers)):
 62        layer_array.append(pretrained_dendrite.dendrite_module.layers[layer_id])
 63        if pretrained_dendrite.dendrite_module.processors == []:
 64            processor_array.append(None)
 65        else:
 66            if not pretrained_dendrite.dendrite_module.processors[layer_id] is None:
 67                pretrained_dendrite.dendrite_module.processors[layer_id].pre = (
 68                    pretrained_dendrite.dendrite_module.processors[layer_id].pre_d
 69                )
 70                pretrained_dendrite.dendrite_module.processors[layer_id].post = (
 71                    pretrained_dendrite.dendrite_module.processors[layer_id].post_d
 72                )
 73            processor_array.append(
 74                pretrained_dendrite.dendrite_module.processors[layer_id]
 75            )
 76    layer_array.append(pretrained_dendrite.main_module)
 77    if not pretrained_dendrite.processor is None:
 78        pretrained_dendrite.processor.pre = pretrained_dendrite.processor.post_n1
 79        pretrained_dendrite.processor.post = pretrained_dendrite.processor.post_n2
 80    processor_array.append(pretrained_dendrite.processor)
 81
 82    view_tuple = []
 83    for dim in range(
 84        len(
 85            pretrained_dendrite.dendrite_module.dendrite_values[0].this_output_dimensions
 86        )
 87    ):
 88        if (
 89            dim
 90            == pretrained_dendrite.dendrite_module.dendrite_values[0].this_node_index
 91        ):
 92            view_tuple.append(-1)
 93            continue
 94        view_tuple.append(1)
 95    return PAILayer(
 96        nn.Sequential(*layer_array),
 97        processor_array,
 98        pretrained_dendrite.dendrites_to_top,
 99        pretrained_dendrite.dendrite_module.dendrites_to_dendrites,
100        pretrained_dendrite.this_node_index,
101        pretrained_dendrite.dendrite_module.num_cycles,
102        view_tuple,
103    )
104
105
106def get_pretrained_pai_attr(pretrained_dendrite, member):
107    if pretrained_dendrite is None:
108        return None
109    else:
110        return getattr(pretrained_dendrite, member)
111
112
113def get_pretrained_pai_var(pretrained_dendrite, submodule_id):
114    if pretrained_dendrite is None:
115        return None
116    else:
117        return pretrained_dendrite[submodule_id]
118
119# This optimizes a network recursively from training modules to inference modules
120def optimize_module(net, depth, name_so_far, converted_list):
121    all_members = net.__dir__()
122    if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList):
123        for submodule_id, layer in net.named_children():
124            if type(net.get_submodule(submodule_id)) is PA.PAINeuronModule:
125                if GPA.pc.get_extra_verbose():
126                    print(
127                        "Seq sub is PAI so optimizing: %s" % name_so_far
128                        + "["
129                        + str(submodule_id)
130                        + "]"
131                    )
132                setattr(
133                    net,
134                    submodule_id,
135                    convert_to_pai_layer_block(net.get_submodule(submodule_id)),
136                )
137            else:
138                if net != net.get_submodule(submodule_id):
139                    # this currently just always returns false, not sure what it was for
140                    converted_list += [name_so_far + "[" + str(submodule_id) + "]"]
141                    setattr(
142                        net,
143                        submodule_id,
144                        optimize_module(
145                            net.get_submodule(submodule_id),
146                            depth + 1,
147                            name_so_far + "[" + str(submodule_id) + "]",
148                            converted_list,
149                        ),
150                    )
151                else:
152                    if GPA.pc.get_extra_verbose():
153                        print(
154                            "%s is a self pointer so skipping"
155                            % (name_so_far + "[" + str(submodule_id) + "]")
156                        )
157    else:
158        for member in all_members:
159            if isinstance(getattr(type(net), member, None), property):
160                continue
161            try:
162                getattr(net, member, None)
163            except:
164                continue
165            sub_name = name_so_far + "." + member
166            if (
167                sub_name in GPA.pc.get_module_names_to_not_save()
168                or sub_name in converted_list
169            ):
170                if GPA.pc.get_extra_verbose():
171                    print("Skipping %s during save" % sub_name)
172                continue
173            if type(getattr(net, member, None)) is PA.PAINeuronModule:
174                if GPA.pc.get_extra_verbose():
175                    print(
176                        "Sub is in conversion list so initiating optimization for: %s"
177                        % name_so_far
178                        + "."
179                        + member
180                    )
181                setattr(net, member, convert_to_pai_layer_block(getattr(net, member)))
182            elif issubclass(type(getattr(net, member, None)), nn.Module):
183                if net != getattr(net, member):
184                    converted_list += [sub_name]
185                    setattr(
186                        net,
187                        member,
188                        optimize_module(
189                            getattr(net, member),
190                            depth + 1,
191                            sub_name,
192                            converted_list,
193                        ),
194                    )
195                else:
196                    if GPA.pc.get_extra_verbose():
197                        print("%s is a self pointer so skipping" % (sub_name))
198    return net
199
200def blockwise_network(net):
201    return optimize_module(net, 0, "", [])
class PAILayer(torch.nn.modules.module.Module): 
15class PAILayer(nn.Module):
16    def __init__(
17        self,
18        layer_array,
19        processor_array,
20        dendrites_to_top,
21        dendrites_to_dendrites,
22        node_index,
23        num_cycles,
24        view_tuple,
25    ):
26        super(PAILayer, self).__init__()
27        self.layer_array = layer_array
28        self.register_buffer("num_cycles", num_cycles)
29        self.register_buffer("view_tuple", torch.tensor(view_tuple))
30
31        self.processor_array = processor_array
32        if dendrites_to_dendrites:
33            self.skip_weights = dendrites_to_dendrites
34        else:
35            """
36            This will only be the case if there is less than 2 dendrites, in these cases an empty array
37            should still be added so that dendrites_to_top is included at the correct index
38            """
39            self.skip_weights = nn.ParameterList()
40        if dendrites_to_top:
41            self.skip_weights.append(dendrites_to_top[len(dendrites_to_top) - 1])
42        else:
43            self.skip_weights = nn.ParameterList()
44        
45        # Delete skip_weights if it's empty (only 1 layer, no skip connections)
46        if len(self.skip_weights) == 0:
47            delattr(self, 'skip_weights')
48
49        self.node_index = node_index
50        self.internal_nonlinearity = GPA.pc.get_pai_forward_function()

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

PAILayer( layer_array, processor_array, dendrites_to_top, dendrites_to_dendrites, node_index, num_cycles, view_tuple) 
16    def __init__(
17        self,
18        layer_array,
19        processor_array,
20        dendrites_to_top,
21        dendrites_to_dendrites,
22        node_index,
23        num_cycles,
24        view_tuple,
25    ):
26        super(PAILayer, self).__init__()
27        self.layer_array = layer_array
28        self.register_buffer("num_cycles", num_cycles)
29        self.register_buffer("view_tuple", torch.tensor(view_tuple))
30
31        self.processor_array = processor_array
32        if dendrites_to_dendrites:
33            self.skip_weights = dendrites_to_dendrites
34        else:
35            """
36            This will only be the case if there is less than 2 dendrites, in these cases an empty array
37            should still be added so that dendrites_to_top is included at the correct index
38            """
39            self.skip_weights = nn.ParameterList()
40        if dendrites_to_top:
41            self.skip_weights.append(dendrites_to_top[len(dendrites_to_top) - 1])
42        else:
43            self.skip_weights = nn.ParameterList()
44        
45        # Delete skip_weights if it's empty (only 1 layer, no skip connections)
46        if len(self.skip_weights) == 0:
47            delattr(self, 'skip_weights')
48
49        self.node_index = node_index
50        self.internal_nonlinearity = GPA.pc.get_pai_forward_function()

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

layer_array
processor_array
node_index
internal_nonlinearity
def unWrap_params(model): 
52def unWrap_params(model):
53    for p in model.parameters():
54        if "wrapped" in p.__dir__():
55            del p.wrapped
def convert_to_pai_layer_block(pretrained_dendrite): 
 58def convert_to_pai_layer_block(pretrained_dendrite):
 59    unWrap_params(pretrained_dendrite)
 60    layer_array = []
 61    processor_array = []
 62    for layer_id in range(len(pretrained_dendrite.dendrite_module.layers)):
 63        layer_array.append(pretrained_dendrite.dendrite_module.layers[layer_id])
 64        if pretrained_dendrite.dendrite_module.processors == []:
 65            processor_array.append(None)
 66        else:
 67            if not pretrained_dendrite.dendrite_module.processors[layer_id] is None:
 68                pretrained_dendrite.dendrite_module.processors[layer_id].pre = (
 69                    pretrained_dendrite.dendrite_module.processors[layer_id].pre_d
 70                )
 71                pretrained_dendrite.dendrite_module.processors[layer_id].post = (
 72                    pretrained_dendrite.dendrite_module.processors[layer_id].post_d
 73                )
 74            processor_array.append(
 75                pretrained_dendrite.dendrite_module.processors[layer_id]
 76            )
 77    layer_array.append(pretrained_dendrite.main_module)
 78    if not pretrained_dendrite.processor is None:
 79        pretrained_dendrite.processor.pre = pretrained_dendrite.processor.post_n1
 80        pretrained_dendrite.processor.post = pretrained_dendrite.processor.post_n2
 81    processor_array.append(pretrained_dendrite.processor)
 82
 83    view_tuple = []
 84    for dim in range(
 85        len(
 86            pretrained_dendrite.dendrite_module.dendrite_values[0].this_output_dimensions
 87        )
 88    ):
 89        if (
 90            dim
 91            == pretrained_dendrite.dendrite_module.dendrite_values[0].this_node_index
 92        ):
 93            view_tuple.append(-1)
 94            continue
 95        view_tuple.append(1)
 96    return PAILayer(
 97        nn.Sequential(*layer_array),
 98        processor_array,
 99        pretrained_dendrite.dendrites_to_top,
100        pretrained_dendrite.dendrite_module.dendrites_to_dendrites,
101        pretrained_dendrite.this_node_index,
102        pretrained_dendrite.dendrite_module.num_cycles,
103        view_tuple,
104    )
def get_pretrained_pai_attr(pretrained_dendrite, member): 
107def get_pretrained_pai_attr(pretrained_dendrite, member):
108    if pretrained_dendrite is None:
109        return None
110    else:
111        return getattr(pretrained_dendrite, member)
def get_pretrained_pai_var(pretrained_dendrite, submodule_id): 
114def get_pretrained_pai_var(pretrained_dendrite, submodule_id):
115    if pretrained_dendrite is None:
116        return None
117    else:
118        return pretrained_dendrite[submodule_id]
def optimize_module(net, depth, name_so_far, converted_list): 
121def optimize_module(net, depth, name_so_far, converted_list):
122    all_members = net.__dir__()
123    if issubclass(type(net), nn.Sequential) or issubclass(type(net), nn.ModuleList):
124        for submodule_id, layer in net.named_children():
125            if type(net.get_submodule(submodule_id)) is PA.PAINeuronModule:
126                if GPA.pc.get_extra_verbose():
127                    print(
128                        "Seq sub is PAI so optimizing: %s" % name_so_far
129                        + "["
130                        + str(submodule_id)
131                        + "]"
132                    )
133                setattr(
134                    net,
135                    submodule_id,
136                    convert_to_pai_layer_block(net.get_submodule(submodule_id)),
137                )
138            else:
139                if net != net.get_submodule(submodule_id):
140                    # this currently just always returns false, not sure what it was for
141                    converted_list += [name_so_far + "[" + str(submodule_id) + "]"]
142                    setattr(
143                        net,
144                        submodule_id,
145                        optimize_module(
146                            net.get_submodule(submodule_id),
147                            depth + 1,
148                            name_so_far + "[" + str(submodule_id) + "]",
149                            converted_list,
150                        ),
151                    )
152                else:
153                    if GPA.pc.get_extra_verbose():
154                        print(
155                            "%s is a self pointer so skipping"
156                            % (name_so_far + "[" + str(submodule_id) + "]")
157                        )
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            if (
168                sub_name in GPA.pc.get_module_names_to_not_save()
169                or sub_name in converted_list
170            ):
171                if GPA.pc.get_extra_verbose():
172                    print("Skipping %s during save" % sub_name)
173                continue
174            if type(getattr(net, member, None)) is PA.PAINeuronModule:
175                if GPA.pc.get_extra_verbose():
176                    print(
177                        "Sub is in conversion list so initiating optimization for: %s"
178                        % name_so_far
179                        + "."
180                        + member
181                    )
182                setattr(net, member, convert_to_pai_layer_block(getattr(net, member)))
183            elif issubclass(type(getattr(net, member, None)), nn.Module):
184                if net != getattr(net, member):
185                    converted_list += [sub_name]
186                    setattr(
187                        net,
188                        member,
189                        optimize_module(
190                            getattr(net, member),
191                            depth + 1,
192                            sub_name,
193                            converted_list,
194                        ),
195                    )
196                else:
197                    if GPA.pc.get_extra_verbose():
198                        print("%s is a self pointer so skipping" % (sub_name))
199    return net
def blockwise_network(net): 
201def blockwise_network(net):
202    return optimize_module(net, 0, "", [])