Package: torch-compose
ModuleGraphInstallationpipinstalltorch-composeIntroductionThe ModuleGraph project is a Python-based tool that facilitates managing and manipulating a network of PyTorch modules. The tool employs directed graphs, allowing each module to use the output of another as its input, thus accommodating complex data flows within a neural network.ModuleGraph is ideal in cases where typical sequential or parallel module arrangements don't suffice, allowing intricate dependencies between modules and handling data propagation in the correct order.FeaturesDirectedModule: This versatile abstract base class represents a module that receives and produces data. It requires implementing a forward function within its subclass. It can be integrated into a module graph, providing a controlled data flow from one module to another.TheDirectedModuleaccepts and produces data using defined keys, providing fine control over data partitioning and labeling. It can wrap an existing PyTorch module, reference an existing function as the forward method, or be subclassed for custom forward method implementation, offering great flexibility for various deep learning tasks.ModuleGraph: This class defines a graph ofDirectedModules. The data can flow from one module to another according to the topological order of dependencies. TheModuleGraphtakes care of ensuring that the modules are correctly sorted to respect their dependencies.TheModuleGraphclass also provides a visualization methodshow_graph()to visually inspect the graph structure and dependencies using thenetworkxandmatplotliblibraries. This visual representation shows modules as nodes and dependencies as directed edges, labeled with the corresponding keys.UsageTo use the ModuleGraph project, you need to define yourDirectedModules and then pass them as a dictionary to theModuleGraphclass.Below is a simplified example:Example 1: Usingmoduleargument toDirectedModuleclassMyModule(DirectedModule):defforward(self,x):returntorch.relu(x)# create DirectedModulesmod1=DirectedModule(input_keys='input',output_keys='hidden',module=MyModule())mod2=DirectedModule(input_keys='hidden',output_keys='output',module=MyModule())# create ModuleGraphgraph=ModuleGraph({'mod1':mod1,'mod2':mod2})# forward passoutput=graph.forward({'input':torch.randn(10,10)})# visualize the graphgraph.show_graph()DirectedModulecan be integrated withtorchcode in several additional ways.`Example 2: Usingforwardargument toDirectedModule# Defining a function that applies a non-linear activation (ReLU) to its inputdefmy_relu(input_tensor):returntorch.nn.functional.relu(input_tensor)# Create a DirectedModule that wraps this functionrelu_module=DirectedModule(input_keys='input',output_keys='output',forward=my_relu)# Use the module in a graphgraph=ModuleGraph({'relu':relu_module})# forward passoutput=graph.forward({'input':torch.randn(10,10)})In this example, therelu_moduledoes not encapsulate a PyTorchnn.Module, but rather a standalone function that performs an operation using PyTorch functionalities.Example 3: UsingDirectedModulein inheritance# Defining a custom module by inheriting from both nn.Module and DirectedModuleclassMyCustomModule(nn.Module,DirectedModule):def__init__(self,input_keys,output_keys):nn.Module.__init__(self)DirectedModule.__init__(self,input_keys,output_keys)self.linear=nn.Linear(10,10)defforward(self,x):returnself.linear(x)# Create an instance of the custom modulecustom_module=MyCustomModule(input_keys='input',output_keys='output')# Use the module in a graphgraph=ModuleGraph({'custom_module':custom_module})# forward passoutput=graph.forward({'input':torch.randn(10,10)})In this example,MyCustomModuleis a subclass of bothnn.ModuleandDirectedModule. It usesnn.Moduleto define a simple linear layer andDirectedModuleto handle input/output key mapping.Example 4: UsingDirectedModuleas a mixin# Defining a custom module by inheriting from both nn.Module and DirectedModuleclassMyMixinModule(nn.Module):def__init__(self,input_keys,output_keys):super().__init__()self.directed_module=DirectedModule(input_keys=input_keys,output_keys=output_keys,module=self)self.linear=nn.Linear(10,10)defforward(self,x):returnself.linear(x)# Create an instance of the mixin modulemixin_module=MyMixinModule(input_keys='input',output_keys='output')# Use the module in a graphgraph=ModuleGraph({'mixin_module':mixin_module.directed_module})# forward passoutput=graph.forward({'input':torch.randn(10,10)})In this example,MyMixinModuleis a subclass ofnn.Module, and it uses aDirectedModuleas a member to handle input/output key mapping, allowing the same functionality without multiple inheritance.