def __tag_nodes()

in tinynn/graph/tracer.py [0:0]

• 65 lines of code
• 24 McCabe index (conditional complexity)

    def __tag_nodes(self) -> None:
        """Gives the modules and the submodules a unique name"""
        # Tag submodules
        for n, m in self.module.named_modules():
            n = n.replace(".", "_")
            n = n.replace("-", "_")
            n = n.replace("/", "_")
            n = n.replace(";", "_")
            n = 'module_' + n if n.isnumeric() else n
            self.module_unique_name_dict[id(m)] = n

        # Tag the module itself
        self.module_unique_name_dict[id(self.module)] = type(self.module).__name__

        for n, p in self.module.named_parameters():
            n = n.replace(".", "_")
            n = n.replace("-", "_")
            n = n.replace("/", "_")
            n = n.replace(";", "_")
            n = 'param_' + n if n.isnumeric() else n
            self.parameter_unique_name_dict[id(p)] = n

        for n, b in self.module.named_buffers():
            n = n.replace(".", "_")
            n = n.replace("-", "_")
            n = n.replace("/", "_")
            n = n.replace(";", "_")
            n = 'buffer_' + n if n.isnumeric() else n
            self.parameter_unique_name_dict[id(b)] = n

        q = queue.Queue()
        q.put((self.module, ''))

        while not q.empty():
            m, p = q.get()
            if isinstance(m, nn.Module):
                self.module_original_name_dict[id(m)] = p
                for n, c in m.named_children():
                    if isinstance(m, (nn.Sequential, nn.ModuleList)) and n.isnumeric():
                        c_p = f'{p}[{n}]'
                    elif isinstance(m, nn.ModuleDict):
                        c_p = f'{p}["{n}"]'
                    else:
                        if len(p) > 0:
                            if '.' in n or '-' in n:
                                c_p = f'{p}.get_submodule("{n}")'
                            else:
                                c_p = f'{p}.{n}'
                        else:
                            c_p = n
                    q.put((c, c_p))
                for n, c in m.named_parameters(recurse=False):
                    self.parameter_module_dict[id(c)] = id(m)
                    if len(p) > 0:
                        if '.' in n or '-' in n:
                            c_p = f'{p}.get_parameter("{n}")'
                        else:
                            c_p = f'{p}.{n}'
                    else:
                        c_p = n
                    q.put((c, c_p))

                for n, c in m.named_buffers(recurse=False):
                    self.parameter_module_dict[id(c)] = id(m)
                    if len(p) > 0:
                        if '.' in n or '-' in n:
                            c_p = f'{p}.get_buffer("{n}")'
                        else:
                            c_p = f'{p}.{n}'
                    else:
                        c_p = n
                    q.put((c, c_p))
            else:
                self.parameter_original_name_dict[id(m)] = p