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        """Initialize a threaded inference wrapper from an existing PAI module.
 19
 20        Parameters
 21        ----------
 22        original_module : nn.Module
 23            Existing PAI module that provides layers, processors, and buffers.
 24        """
 25        super(PAIModulePyThread, self).__init__()
 26        self.layer_array = original_module.layer_array
 27        self.processor_array = original_module.processor_array
 28        # Remove the unused first index (skip_weights[0] is never used)
 29        if hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) > 1:
 30            self.skip_weights = original_module.skip_weights[1:]
 31        elif hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) == 1:
 32            # Only one element, don't create skip_weights
 33            pass
 34        self.register_buffer("node_index", original_module.node_index.clone().detach())
 35        self.register_buffer("num_cycles", original_module.num_cycles)
 36        self.register_buffer("view_tuple", original_module.view_tuple)
 37
 38    def process_and_forward(self, *args2, **kwargs2):
 39        """Run one dendrite layer forward pass and store its output.
 40
 41        Parameters
 42        ----------
 43        *args2 : tuple
 44            Positional values where the first two entries are layer index and
 45            shared output buffer.
 46        **kwargs2 : dict
 47            Keyword arguments forwarded to the wrapped layer.
 48
 49        Returns
 50        -------
 51        None
 52            This function does not return a value.
 53        """
 54        c = args2[0]
 55        dendrite_outs = args2[1]
 56        args2 = args2[2:]
 57        if self.processor_array[c] != None:
 58            args2, kwargs2 = self.processor_array[c].pre(*args2, **kwargs2)
 59        out_values = self.layer_array[c](*args2, **kwargs2)
 60        if self.processor_array[c] != None:
 61            out = self.processor_array[c].post(out_values)
 62        else:
 63            out = out_values
 64        dendrite_outs[c] = out
 65
 66    def process_and_pre(self, *args, **kwargs):
 67        """Run the final layer pre-pass used before skip accumulation.
 68
 69        Parameters
 70        ----------
 71        *args : tuple
 72            Positional values where the first entry is the shared output
 73            buffer and the remaining values are layer inputs.
 74        **kwargs : dict
 75            Keyword arguments forwarded to the wrapped layer.
 76
 77        Returns
 78        -------
 79        None
 80            This function does not return a value.
 81        """
 82        dendrite_outs = args[0]
 83        args = args[1:]
 84        out = self.layer_array[-1].forward(*args, **kwargs)
 85        if not self.processor_array[-1] is None:
 86            out = self.processor_array[-1].pre(out)
 87        dendrite_outs[len(self.layer_array) - 1] = out
 88
 89    def forward(self, *args, **kwargs):
 90        """Compute module output with optional threaded dendrite evaluation.
 91
 92        Parameters
 93        ----------
 94        *args : tuple
 95            Positional arguments passed into each wrapped layer.
 96        **kwargs : dict
 97            Keyword arguments passed into each wrapped layer.
 98
 99        Returns
100        -------
101        Any
102            Final module output after skip connections and post-processing.
103        """
104        # this is currently false anyway, just remove the doing multi idea
105        doing_multi = doing_threading
106        dendrite_outs = [None] * len(self.layer_array)
107        threads = {}
108        for c in range(0, len(self.layer_array) - 1):
109            args2, kwargs2 = args, kwargs
110            if doing_multi:
111                threads[c] = Thread(
112                    target=self.process_and_forward,
113                    args=(c, dendrite_outs, *args),
114                    kwargs=kwargs,
115                )
116            else:
117                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
118        if doing_multi:
119            threads[len(self.layer_array) - 1] = Thread(
120                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
121            )
122        else:
123            self.process_and_pre(dendrite_outs, *args, **kwargs)
124        if doing_multi:
125            for i in range(len(dendrite_outs)):
126                threads[i].start()
127            for i in range(len(dendrite_outs)):
128                threads[i].join()
129        for out_index in range(0, len(self.layer_array)):
130            current_out = dendrite_outs[out_index]
131            if len(self.layer_array) > 1 and hasattr(self, 'skip_weights'):
132                for in_index in range(0, out_index):
133                    # Use out_index - 1 because skip_weights[0] was removed
134                    skip_weight = self.skip_weights[out_index - 1][in_index, :]
135                    # Use cached Python tuple instead of .tolist() during forward
136                    skip_weight = skip_weight.view(self.view_tuple.tolist())
137                    current_out = current_out + (
138                        skip_weight.to(current_out.device)
139                        * dendrite_outs[in_index]
140                    )
141                if out_index < len(self.layer_array) - 1:
142                    current_out = GPA.pc.get_pai_forward_function()(current_out)
143            dendrite_outs[out_index] = current_out
144        if not self.processor_array[-1] is None:
145            current_out = self.processor_array[-1].post(current_out)
146        return current_out
147
148
149def get_pretrained_pai_attr(pretrained_dendrite, member):
150    """Safely get an attribute from a possibly missing module.
151
152    Parameters
153    ----------
154    pretrained_dendrite : nn.Module or None
155        Source module that may be ``None``.
156    member : str
157        Attribute name to retrieve.
158
159    Returns
160    -------
161    Any
162        Requested attribute value, or ``None`` if source module is ``None``.
163    """
164    if pretrained_dendrite is None:
165        return None
166    else:
167        return getattr(pretrained_dendrite, member)
168
169
170def get_pretrained_pai_var(pretrained_dendrite, submodule_id):
171    """Safely get a named child module from a possibly missing module.
172
173    Parameters
174    ----------
175    pretrained_dendrite : nn.Module or None
176        Source module that may be ``None``.
177    submodule_id : str
178        Submodule identifier passed to ``get_submodule``.
179
180    Returns
181    -------
182    nn.Module or None
183        Retrieved submodule, or ``None`` when source module is ``None``.
184    """
185    if pretrained_dendrite is None:
186        return None
187    else:
188        return pretrained_dendrite.get_submodule(submodule_id)
189
190ModuleType = PAIModulePyThread
191doing_threading = False
192
193def make_module(module):
194    """Create the configured wrapper module type for a module.
195
196    Parameters
197    ----------
198    module : nn.Module
199        Module to wrap.
200
201    Returns
202    -------
203    nn.Module
204        Wrapped module instance.
205    """
206    return ModuleType(module)
207
208# This Refreshes a PAI network with the PyThread Module
209def refresh_pai(net, depth, name_so_far, converted_list):
210    """Recursively replace PAILayer instances with threaded inference wrappers.
211
212    Parameters
213    ----------
214    net : nn.Module
215        Module tree to update.
216    depth : int
217        Current recursion depth.
218    name_so_far : str
219        Dotted/indexed path to the current module.
220    converted_list : list
221        Mutable list of module names already visited.
222
223    Returns
224    -------
225    nn.Module
226        Updated module tree.
227    """
228    if GPA.pc.get_extra_verbose():
229        print("CL calling convert on %s depth %d" % (net, depth))
230        print(
231            "CL calling convert on %s: %s, depth %d"
232            % (name_so_far, type(net).__name__, depth)
233        )
234    if type(net) is ModuleType:
235        if GPA.pc.get_extra_verbose():
236            print(
237                "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"
238            )
239        return net
240    all_members = net.__dir__()
241    if (
242        issubclass(type(net), nn.Sequential)
243        or issubclass(type(net), nn.ModuleList)
244        or issubclass(type(net), list)
245    ):
246        for submodule_id, layer in net.named_children():
247            if net != net.get_submodule(submodule_id):
248                converted_list += [name_so_far + "[" + str(submodule_id) + "]"]
249                setattr(
250                    net,
251                    submodule_id,
252                    refresh_pai(
253                        net.get_submodule(submodule_id),
254                        depth + 1,
255                        name_so_far + "[" + str(submodule_id) + "]",
256                        converted_list,
257                    ),
258                )
259            if type(net.get_submodule(submodule_id)).__name__ == "PAILayer":
260                setattr(
261                    net,
262                    submodule_id,
263                    make_module(get_pretrained_pai_var(net, submodule_id)),
264                )
265    elif type(net) in GPA.pc.get_modules_to_track():
266        return net
267    else:
268        for member in all_members:
269            if isinstance(getattr(type(net), member, None), property):
270                continue
271            try:
272                getattr(net, member, None)
273            except:
274                continue
275            sub_name = name_so_far + "." + member
276
277            if member == "device" or member == "dtype":
278                continue
279            if sub_name in GPA.pc.get_module_names_to_not_save():
280                continue
281            if name_so_far == "":
282                if (
283                    sub_name in GPA.pc.get_module_names_to_not_save()
284                    or sub_name in converted_list
285                ):
286                    if GPA.pc.get_extra_verbose():
287                        print("Skipping %s during save" % sub_name)
288                    continue
289
290            if (
291                issubclass(type(getattr(net, member, None)), nn.Module)
292                or member == "layer_array"
293            ):
294                converted_list += [sub_name]
295                if net != getattr(net, member):
296                    setattr(
297                        net,
298                        member,
299                        refresh_pai(
300                            getattr(net, member), depth + 1, sub_name, converted_list
301                        ),
302                    )
303            if type(getattr(net, member, None)).__name__ == "PAILayer":
304                setattr(net, member, make_module(get_pretrained_pai_attr(net, member)))
305    if type(net).__name__ == "PAILayer":
306        net = make_module(net)
307    return net
308
309def refresh_net(pretrained_dendrite):
310    """Refresh a network by converting eligible modules to threaded wrappers.
311
312    Parameters
313    ----------
314    pretrained_dendrite : nn.Module
315        Network or module tree to refresh.
316
317    Returns
318    -------
319    nn.Module
320        Refreshed network.
321    """
322
323    net = refresh_pai(pretrained_dendrite, 0, "", [])
324    return net
doing_threading = False
class PAIModulePyThread(torch.nn.modules.module.Module): 
 17class PAIModulePyThread(nn.Module):
 18    def __init__(self, original_module):
 19        """Initialize a threaded inference wrapper from an existing PAI module.
 20
 21        Parameters
 22        ----------
 23        original_module : nn.Module
 24            Existing PAI module that provides layers, processors, and buffers.
 25        """
 26        super(PAIModulePyThread, self).__init__()
 27        self.layer_array = original_module.layer_array
 28        self.processor_array = original_module.processor_array
 29        # Remove the unused first index (skip_weights[0] is never used)
 30        if hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) > 1:
 31            self.skip_weights = original_module.skip_weights[1:]
 32        elif hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) == 1:
 33            # Only one element, don't create skip_weights
 34            pass
 35        self.register_buffer("node_index", original_module.node_index.clone().detach())
 36        self.register_buffer("num_cycles", original_module.num_cycles)
 37        self.register_buffer("view_tuple", original_module.view_tuple)
 38
 39    def process_and_forward(self, *args2, **kwargs2):
 40        """Run one dendrite layer forward pass and store its output.
 41
 42        Parameters
 43        ----------
 44        *args2 : tuple
 45            Positional values where the first two entries are layer index and
 46            shared output buffer.
 47        **kwargs2 : dict
 48            Keyword arguments forwarded to the wrapped layer.
 49
 50        Returns
 51        -------
 52        None
 53            This function does not return a value.
 54        """
 55        c = args2[0]
 56        dendrite_outs = args2[1]
 57        args2 = args2[2:]
 58        if self.processor_array[c] != None:
 59            args2, kwargs2 = self.processor_array[c].pre(*args2, **kwargs2)
 60        out_values = self.layer_array[c](*args2, **kwargs2)
 61        if self.processor_array[c] != None:
 62            out = self.processor_array[c].post(out_values)
 63        else:
 64            out = out_values
 65        dendrite_outs[c] = out
 66
 67    def process_and_pre(self, *args, **kwargs):
 68        """Run the final layer pre-pass used before skip accumulation.
 69
 70        Parameters
 71        ----------
 72        *args : tuple
 73            Positional values where the first entry is the shared output
 74            buffer and the remaining values are layer inputs.
 75        **kwargs : dict
 76            Keyword arguments forwarded to the wrapped layer.
 77
 78        Returns
 79        -------
 80        None
 81            This function does not return a value.
 82        """
 83        dendrite_outs = args[0]
 84        args = args[1:]
 85        out = self.layer_array[-1].forward(*args, **kwargs)
 86        if not self.processor_array[-1] is None:
 87            out = self.processor_array[-1].pre(out)
 88        dendrite_outs[len(self.layer_array) - 1] = out
 89
 90    def forward(self, *args, **kwargs):
 91        """Compute module output with optional threaded dendrite evaluation.
 92
 93        Parameters
 94        ----------
 95        *args : tuple
 96            Positional arguments passed into each wrapped layer.
 97        **kwargs : dict
 98            Keyword arguments passed into each wrapped layer.
 99
100        Returns
101        -------
102        Any
103            Final module output after skip connections and post-processing.
104        """
105        # this is currently false anyway, just remove the doing multi idea
106        doing_multi = doing_threading
107        dendrite_outs = [None] * len(self.layer_array)
108        threads = {}
109        for c in range(0, len(self.layer_array) - 1):
110            args2, kwargs2 = args, kwargs
111            if doing_multi:
112                threads[c] = Thread(
113                    target=self.process_and_forward,
114                    args=(c, dendrite_outs, *args),
115                    kwargs=kwargs,
116                )
117            else:
118                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
119        if doing_multi:
120            threads[len(self.layer_array) - 1] = Thread(
121                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
122            )
123        else:
124            self.process_and_pre(dendrite_outs, *args, **kwargs)
125        if doing_multi:
126            for i in range(len(dendrite_outs)):
127                threads[i].start()
128            for i in range(len(dendrite_outs)):
129                threads[i].join()
130        for out_index in range(0, len(self.layer_array)):
131            current_out = dendrite_outs[out_index]
132            if len(self.layer_array) > 1 and hasattr(self, 'skip_weights'):
133                for in_index in range(0, out_index):
134                    # Use out_index - 1 because skip_weights[0] was removed
135                    skip_weight = self.skip_weights[out_index - 1][in_index, :]
136                    # Use cached Python tuple instead of .tolist() during forward
137                    skip_weight = skip_weight.view(self.view_tuple.tolist())
138                    current_out = current_out + (
139                        skip_weight.to(current_out.device)
140                        * dendrite_outs[in_index]
141                    )
142                if out_index < len(self.layer_array) - 1:
143                    current_out = GPA.pc.get_pai_forward_function()(current_out)
144            dendrite_outs[out_index] = current_out
145        if not self.processor_array[-1] is None:
146            current_out = self.processor_array[-1].post(current_out)
147        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        """Initialize a threaded inference wrapper from an existing PAI module.
20
21        Parameters
22        ----------
23        original_module : nn.Module
24            Existing PAI module that provides layers, processors, and buffers.
25        """
26        super(PAIModulePyThread, self).__init__()
27        self.layer_array = original_module.layer_array
28        self.processor_array = original_module.processor_array
29        # Remove the unused first index (skip_weights[0] is never used)
30        if hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) > 1:
31            self.skip_weights = original_module.skip_weights[1:]
32        elif hasattr(original_module, 'skip_weights') and len(original_module.skip_weights) == 1:
33            # Only one element, don't create skip_weights
34            pass
35        self.register_buffer("node_index", original_module.node_index.clone().detach())
36        self.register_buffer("num_cycles", original_module.num_cycles)
37        self.register_buffer("view_tuple", original_module.view_tuple)

Initialize a threaded inference wrapper from an existing PAI module.

Parameters
  • original_module (nn.Module): Existing PAI module that provides layers, processors, and buffers.
layer_array
processor_array
def process_and_forward(self, *args2, **kwargs2): 
39    def process_and_forward(self, *args2, **kwargs2):
40        """Run one dendrite layer forward pass and store its output.
41
42        Parameters
43        ----------
44        *args2 : tuple
45            Positional values where the first two entries are layer index and
46            shared output buffer.
47        **kwargs2 : dict
48            Keyword arguments forwarded to the wrapped layer.
49
50        Returns
51        -------
52        None
53            This function does not return a value.
54        """
55        c = args2[0]
56        dendrite_outs = args2[1]
57        args2 = args2[2:]
58        if self.processor_array[c] != None:
59            args2, kwargs2 = self.processor_array[c].pre(*args2, **kwargs2)
60        out_values = self.layer_array[c](*args2, **kwargs2)
61        if self.processor_array[c] != None:
62            out = self.processor_array[c].post(out_values)
63        else:
64            out = out_values
65        dendrite_outs[c] = out

Run one dendrite layer forward pass and store its output.

Parameters
  • *args2 (tuple): Positional values where the first two entries are layer index and shared output buffer.
  • **kwargs2 (dict): Keyword arguments forwarded to the wrapped layer.
Returns
  • None: This function does not return a value.
def process_and_pre(self, *args, **kwargs): 
67    def process_and_pre(self, *args, **kwargs):
68        """Run the final layer pre-pass used before skip accumulation.
69
70        Parameters
71        ----------
72        *args : tuple
73            Positional values where the first entry is the shared output
74            buffer and the remaining values are layer inputs.
75        **kwargs : dict
76            Keyword arguments forwarded to the wrapped layer.
77
78        Returns
79        -------
80        None
81            This function does not return a value.
82        """
83        dendrite_outs = args[0]
84        args = args[1:]
85        out = self.layer_array[-1].forward(*args, **kwargs)
86        if not self.processor_array[-1] is None:
87            out = self.processor_array[-1].pre(out)
88        dendrite_outs[len(self.layer_array) - 1] = out

Run the final layer pre-pass used before skip accumulation.

Parameters
  • *args (tuple): Positional values where the first entry is the shared output buffer and the remaining values are layer inputs.
  • **kwargs (dict): Keyword arguments forwarded to the wrapped layer.
Returns
  • None: This function does not return a value.
def forward(self, *args, **kwargs): 
 90    def forward(self, *args, **kwargs):
 91        """Compute module output with optional threaded dendrite evaluation.
 92
 93        Parameters
 94        ----------
 95        *args : tuple
 96            Positional arguments passed into each wrapped layer.
 97        **kwargs : dict
 98            Keyword arguments passed into each wrapped layer.
 99
100        Returns
101        -------
102        Any
103            Final module output after skip connections and post-processing.
104        """
105        # this is currently false anyway, just remove the doing multi idea
106        doing_multi = doing_threading
107        dendrite_outs = [None] * len(self.layer_array)
108        threads = {}
109        for c in range(0, len(self.layer_array) - 1):
110            args2, kwargs2 = args, kwargs
111            if doing_multi:
112                threads[c] = Thread(
113                    target=self.process_and_forward,
114                    args=(c, dendrite_outs, *args),
115                    kwargs=kwargs,
116                )
117            else:
118                self.process_and_forward(c, dendrite_outs, *args2, **kwargs2)
119        if doing_multi:
120            threads[len(self.layer_array) - 1] = Thread(
121                target=self.process_and_pre, args=(dendrite_outs, *args), kwargs=kwargs
122            )
123        else:
124            self.process_and_pre(dendrite_outs, *args, **kwargs)
125        if doing_multi:
126            for i in range(len(dendrite_outs)):
127                threads[i].start()
128            for i in range(len(dendrite_outs)):
129                threads[i].join()
130        for out_index in range(0, len(self.layer_array)):
131            current_out = dendrite_outs[out_index]
132            if len(self.layer_array) > 1 and hasattr(self, 'skip_weights'):
133                for in_index in range(0, out_index):
134                    # Use out_index - 1 because skip_weights[0] was removed
135                    skip_weight = self.skip_weights[out_index - 1][in_index, :]
136                    # Use cached Python tuple instead of .tolist() during forward
137                    skip_weight = skip_weight.view(self.view_tuple.tolist())
138                    current_out = current_out + (
139                        skip_weight.to(current_out.device)
140                        * dendrite_outs[in_index]
141                    )
142                if out_index < len(self.layer_array) - 1:
143                    current_out = GPA.pc.get_pai_forward_function()(current_out)
144            dendrite_outs[out_index] = current_out
145        if not self.processor_array[-1] is None:
146            current_out = self.processor_array[-1].post(current_out)
147        return current_out

Compute module output with optional threaded dendrite evaluation.

Parameters
  • *args (tuple): Positional arguments passed into each wrapped layer.
  • **kwargs (dict): Keyword arguments passed into each wrapped layer.
Returns
  • Any: Final module output after skip connections and post-processing.
def get_pretrained_pai_attr(pretrained_dendrite, member): 
150def get_pretrained_pai_attr(pretrained_dendrite, member):
151    """Safely get an attribute from a possibly missing module.
152
153    Parameters
154    ----------
155    pretrained_dendrite : nn.Module or None
156        Source module that may be ``None``.
157    member : str
158        Attribute name to retrieve.
159
160    Returns
161    -------
162    Any
163        Requested attribute value, or ``None`` if source module is ``None``.
164    """
165    if pretrained_dendrite is None:
166        return None
167    else:
168        return getattr(pretrained_dendrite, member)

Safely get an attribute from a possibly missing module.

Parameters
  • pretrained_dendrite (nn.Module or None): Source module that may be None.
  • member (str): Attribute name to retrieve.
Returns
  • Any: Requested attribute value, or None if source module is None.
def get_pretrained_pai_var(pretrained_dendrite, submodule_id): 
171def get_pretrained_pai_var(pretrained_dendrite, submodule_id):
172    """Safely get a named child module from a possibly missing module.
173
174    Parameters
175    ----------
176    pretrained_dendrite : nn.Module or None
177        Source module that may be ``None``.
178    submodule_id : str
179        Submodule identifier passed to ``get_submodule``.
180
181    Returns
182    -------
183    nn.Module or None
184        Retrieved submodule, or ``None`` when source module is ``None``.
185    """
186    if pretrained_dendrite is None:
187        return None
188    else:
189        return pretrained_dendrite.get_submodule(submodule_id)

Safely get a named child module from a possibly missing module.

Parameters
  • pretrained_dendrite (nn.Module or None): Source module that may be None.
  • submodule_id (str): Submodule identifier passed to get_submodule.
Returns
  • nn.Module or None: Retrieved submodule, or None when source module is None.
ModuleType = <class 'PAIModulePyThread'>
def make_module(module): 
194def make_module(module):
195    """Create the configured wrapper module type for a module.
196
197    Parameters
198    ----------
199    module : nn.Module
200        Module to wrap.
201
202    Returns
203    -------
204    nn.Module
205        Wrapped module instance.
206    """
207    return ModuleType(module)

Create the configured wrapper module type for a module.

Parameters
  • module (nn.Module): Module to wrap.
Returns
  • nn.Module: Wrapped module instance.
def refresh_pai(net, depth, name_so_far, converted_list): 
210def refresh_pai(net, depth, name_so_far, converted_list):
211    """Recursively replace PAILayer instances with threaded inference wrappers.
212
213    Parameters
214    ----------
215    net : nn.Module
216        Module tree to update.
217    depth : int
218        Current recursion depth.
219    name_so_far : str
220        Dotted/indexed path to the current module.
221    converted_list : list
222        Mutable list of module names already visited.
223
224    Returns
225    -------
226    nn.Module
227        Updated module tree.
228    """
229    if GPA.pc.get_extra_verbose():
230        print("CL calling convert on %s depth %d" % (net, depth))
231        print(
232            "CL calling convert on %s: %s, depth %d"
233            % (name_so_far, type(net).__name__, depth)
234        )
235    if type(net) is ModuleType:
236        if GPA.pc.get_extra_verbose():
237            print(
238                "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"
239            )
240        return net
241    all_members = net.__dir__()
242    if (
243        issubclass(type(net), nn.Sequential)
244        or issubclass(type(net), nn.ModuleList)
245        or issubclass(type(net), list)
246    ):
247        for submodule_id, layer in net.named_children():
248            if net != net.get_submodule(submodule_id):
249                converted_list += [name_so_far + "[" + str(submodule_id) + "]"]
250                setattr(
251                    net,
252                    submodule_id,
253                    refresh_pai(
254                        net.get_submodule(submodule_id),
255                        depth + 1,
256                        name_so_far + "[" + str(submodule_id) + "]",
257                        converted_list,
258                    ),
259                )
260            if type(net.get_submodule(submodule_id)).__name__ == "PAILayer":
261                setattr(
262                    net,
263                    submodule_id,
264                    make_module(get_pretrained_pai_var(net, submodule_id)),
265                )
266    elif type(net) in GPA.pc.get_modules_to_track():
267        return net
268    else:
269        for member in all_members:
270            if isinstance(getattr(type(net), member, None), property):
271                continue
272            try:
273                getattr(net, member, None)
274            except:
275                continue
276            sub_name = name_so_far + "." + member
277
278            if member == "device" or member == "dtype":
279                continue
280            if sub_name in GPA.pc.get_module_names_to_not_save():
281                continue
282            if name_so_far == "":
283                if (
284                    sub_name in GPA.pc.get_module_names_to_not_save()
285                    or sub_name in converted_list
286                ):
287                    if GPA.pc.get_extra_verbose():
288                        print("Skipping %s during save" % sub_name)
289                    continue
290
291            if (
292                issubclass(type(getattr(net, member, None)), nn.Module)
293                or member == "layer_array"
294            ):
295                converted_list += [sub_name]
296                if net != getattr(net, member):
297                    setattr(
298                        net,
299                        member,
300                        refresh_pai(
301                            getattr(net, member), depth + 1, sub_name, converted_list
302                        ),
303                    )
304            if type(getattr(net, member, None)).__name__ == "PAILayer":
305                setattr(net, member, make_module(get_pretrained_pai_attr(net, member)))
306    if type(net).__name__ == "PAILayer":
307        net = make_module(net)
308    return net

Recursively replace PAILayer instances with threaded inference wrappers.

Parameters
  • net (nn.Module): Module tree to update.
  • depth (int): Current recursion depth.
  • name_so_far (str): Dotted/indexed path to the current module.
  • converted_list (list): Mutable list of module names already visited.
Returns
  • nn.Module: Updated module tree.
def refresh_net(pretrained_dendrite): 
310def refresh_net(pretrained_dendrite):
311    """Refresh a network by converting eligible modules to threaded wrappers.
312
313    Parameters
314    ----------
315    pretrained_dendrite : nn.Module
316        Network or module tree to refresh.
317
318    Returns
319    -------
320    nn.Module
321        Refreshed network.
322    """
323
324    net = refresh_pai(pretrained_dendrite, 0, "", [])
325    return net

Refresh a network by converting eligible modules to threaded wrappers.

Parameters
  • pretrained_dendrite (nn.Module): Network or module tree to refresh.
Returns
  • nn.Module: Refreshed network.