Package: hyper-nn
hyper-nn -- Easy Hypernetworks in Pytorch and FlaxNote: This library is experimental and currently under development - the flax implementations in particular are far from perfect and can be improved. If you have any suggestions on how to improve this library, please open a github issue or feel free to reach out directly!hyper-nngives users with the ability to create easily customizableHypernetworksfor almost any generictorch.nn.ModulefromPytorchandflax.linen.ModulefromFlax. Our Hypernetwork objects are alsotorch.nn.Modulesandflax.linen.Modules, allowing for easy integration with existing systems. For Pytorch, we make use of the amazing libraryfunctorchGenerating Policy Weights for Lunar LanderDynamic Weights for each character in a name generatorInstallhyper-nntested on python 3.8+Installing with pip$pipinstallhyper-nnInstalling from source$gitclonegit@github.com:shyamsn97/hyper-nn.git
$cdhyper-nn
$pythonsetup.pyinstallFor gpu functionality with Jax, you will need to follow the instructionshereWhat are Hypernetworks?Hypernetworks, simply put, are neural networks that generate parameters for another neural network. They can be incredibly powerful, being able to represent large networks while using only a fraction of their parameters.Hypernetworks generally come in two variants, static or dynamic. Static Hypernetworks have a fixed or learned embedding and weight generator that outputs the target networks’ weights deterministically. Dynamic Hypernetworks instead receive inputs and use them to generate dynamic weights.`hyper-nn` allows you to design Hypernetworks with flexibility and ease, all you have to do is implement the `generate_params` method, which outputs a parameter vector. We also include basic versions of this, composed of two linear components:embedding_modulethat holds information about layers(s) in the target network, or more generally a chunk of the target networks weights. By default this outputs a matrix of sizenum_embeddings x embedding_dimweight_generator, which takes in the embedding and outputs a flat parameter vector for the target network. By default this module outputs chunks in the size ofweight_chunk_dim, which is calculated automatically asnum_target_parameters // num_embeddings.Bothembedding_moduleandweight_generatorare represented astorch.nn.Moduleandflax.linen.Moduleobjects. aModulecan be passed in ascustom_embedding_moduleorcustom_weight_generator, or it can be defined in the methodsmake_embedding_moduleormake_weight_generator.Thegenerate_paramsmethod feeds the output fromembedding_moduleintoweight_generatorto output the target parameters.Theforwardmethod takes in a list of inputs and uses the generated parameters to calculate the output. This method acts as the main method for both jax and torch hypernetworksTorch Hypernetwork...defmake_embedding_module(self)->nn.Module:returnnn.Embedding(self.num_embeddings,self.embedding_dim)defmake_weight_generator(self)->nn.Module:returnnn.Linear(self.embedding_dim,self.weight_chunk_dim)defgenerate_params(self,*args,**kwargs)->Tuple[torch.Tensor,Dict[str,Any]]:embedding=self.embedding_module(torch.arange(self.num_embeddings,device=self.device))generated_params=self.weight_generator(embedding).view(-1)returngenerated_params,{"embedding":embedding}deftarget_forward(self,*args,generated_params:torch.Tensor,assert_parameter_shapes:bool=True,**kwargs,)->torch.Tensor:ifassert_parameter_shapes:self.assert_parameter_shapes(generated_params)returnself.target_network(generated_params,*args,**kwargs)defforward(self,*args,generated_params:Optional[torch.Tensor]=None,has_aux:bool=False,assert_parameter_shapes:bool=True,generate_params_kwargs:Dict[str,Any]={},**kwargs,):"""Main method for creating / using generated parameters and passing in input into the target networkArgs:generated_params (Optional[torch.Tensor], optional): Generated parameters of the target network. If not provided, the hypernetwork will generate the parameters. Defaults to None.has_aux (bool, optional): If True, return the auxiliary output from generate_params method. Defaults to False.assert_parameter_shapes (bool, optional): If True, raise an error if generated_params does not have shape (num_target_parameters,). Defaults to True.generate_params_kwargs (Dict[str, Any], optional): kwargs to be passed to generate_params method*args, *kwargs, arguments to be passed into the target network (also gets passed into generate_params)Returns:output (torch.Tensor) | (torch.Tensor, Dict[str, torch.Tensor]): returns output from target network and optionally auxiliary output."""aux_output={}ifgenerated_paramsisNone:generated_params,aux_output=self.generate_params(*args,**kwargs,**generate_params_kwargs)ifhas_aux:return(self.target_forward(*args,generated_params=generated_params,assert_parameter_shapes=assert_parameter_shapes,**kwargs,),generated_params,aux_output,)returnself.target_forward(*args,generated_params=generated_params,assert_parameter_shapes=assert_parameter_shapes,**kwargs,)...Flax Hypernetwork...defmake_embedding_module(self):returnnn.Embed(self.num_embeddings,self.embedding_dim,embedding_init=jax.nn.initializers.uniform(),)defmake_weight_generator(self):returnnn.Dense(self.weight_chunk_dim)defgenerate_params(self,*args,**kwargs)->Tuple[jnp.array,Dict[str,Any]]:embedding=self.embedding_module(jnp.arange(0,self.num_embeddings))generated_params=self.weight_generator(embedding).reshape(-1)returngenerated_params,{"embedding":embedding}deftarget_forward(self,*args,generated_params:jnp.array,assert_parameter_shapes:bool=True,**kwargs,)->jnp.array:ifassert_parameter_shapes:self.assert_parameter_shapes(generated_params)param_tree=create_param_tree(generated_params,self.target_weight_shapes,self.target_treedef)returnself.target_network.apply(param_tree,*args,**kwargs)defforward(self,*args,generated_params:Optional[jnp.array]=None,has_aux:bool=False,assert_parameter_shapes:bool=True,generate_params_kwargs:Dict[str,Any]={},**kwargs,)->Tuple[jnp.array,List[jnp.array]]:"""Main method for creating / using generated parameters and passing in input into the target networkArgs:generated_params (Optional[jnp.array], optional): Generated parameters of the target network. If not provided, the hypernetwork will generate the parameters. Defaults to None.has_aux (bool, optional): If True, return the auxiliary output from generate_params method. Defaults to False.assert_parameter_shapes (bool, optional): If True, raise an error if generated_params does not have shape (num_target_parameters,). Defaults to True.generate_params_kwargs (Dict[str, Any], optional): kwargs to be passed to generate_params methodReturns:output (torch.Tensor) | (jnp.array, Dict[str, jnp.array]): returns output from target network and optionally auxiliary output."""aux_output={}ifgenerated_paramsisNone:generated_params,aux_output=self.generate_params(*args,**kwargs,**generate_params_kwargs)ifhas_aux:return(self.target_forward(*args,generated_params=generated_params,assert_parameter_shapes=assert_parameter_shapes,**kwargs,),generated_params,aux_output,)returnself.target_forward(*args,generated_params=generated_params,assert_parameter_shapes=assert_parameter_shapes,**kwargs,)...Quick Usagefor detailed examples seenotebooksGenerating weights for a CNN on MNISTLunar Lander Reinforce (Vanilla Policy Gradient)Dynamic Hypernetworks for name generationThe main classes to use areTorchHyperNetworkandJaxHyperNetworkand those that inherit them. Instead of constructing them directly, use thefrom_targetmethod, shown below. After this you can use the hypernetwork exactly like any othernn.Module!hyper-nnalso makes it easy to create Dynamic Hypernetworks that use inputs to create target weights. Basic implementations (both < 100 lines) are provided withJaxDynamicHyperNetworkandTorchDynamicHyperNetwork, which use an rnn and current input to generate weights.To create hypernetworks, its easier to use thefrom_targetmethod instead of instantiating it directly because some parameters are calculated automatically for you.Pytorchimporttorchimporttorch.nnasnnfromhypernn.torchimportTorchHyperNetwork,TorchLinearHyperNetwork,TorchDynamicHyperNetwork# any moduletarget_network=nn.Sequential(nn.Linear(32,64),nn.ReLU(),nn.Linear(64,32))EMBEDDING_DIM=4NUM_EMBEDDINGS=32hypernetwork=TorchLinearHyperNetwork(target_network=target_network,embedding_dim=EMBEDDING_DIM,num_embeddings=NUM_EMBEDDINGS)# now we can use the hypernetwork like any other nn.Moduleinp=torch.zeros((1,32))# by default we only output what we'd expect from the target networkoutput=hypernetwork(inp)# return aux_outputoutput,generated_params,aux_output=hypernetwork(inp,has_aux=True)# generate params separatelygenerated_params,aux_output=hypernetwork.generate_params()output=hypernetwork(inp,generated_params=generated_params)### Dynamic Hypernetworkdynamic_hypernetwork=TorchDynamicHyperNetwork(input_dim=32,target_network=target_network,embedding_dim=EMBEDDING_DIM,num_embeddings=NUM_EMBEDDINGS)output=dynamic_hypernetwork(inp,generate_params_kwargs=dict(x=inp))# by default we only output what we'd expect from the target networkoutput=dynamic_hypernetwork(inp,generate_params_kwargs=dict(x=inp,hidden_state=torch.zeros((1,32))))Jaximportflax.linenasnnimportjax.numpyasjnpfromjaximportrandomfromhypernn.jaximportJaxHyperNetwork,JaxLinearHyperNetwork,JaxDynamicHyperNetwork# any moduletarget_network=nn.Sequential([nn.Dense(64),nn.relu,nn.Dense(32)])EMBEDDING_DIM=4NUM_EMBEDDINGS=32hypernetwork=JaxLinearHyperNetwork.from_target(target_network=target_network,embedding_dim=EMBEDDING_DIM,num_embeddings=NUM_EMBEDDINGS,inputs=jnp.zeros((1,32))# jax needs this to initialize target weights)# now we can use the hypernetwork like any other nn.Moduleinp=jnp.zeros((1,32))key=random.PRNGKey(0)hypernetwork_params=hypernetwork.init(key,inp)# flax needs to initialize hypernetwork parameters first# by default we only output what we'd expect from the target networkoutput=hypernetwork.apply(hypernetwork_params,inp)# return aux_outputoutput,generated_params,aux_output=hypernetwork.apply(hypernetwork_params,inp,has_aux=True)# generate params separatelygenerated_params,aux_output=hypernetwork.apply(hypernetwork_params,method=hypernetwork.generate_params)output=hypernetwork.apply(hypernetwork_params,inp,generated_params=generated_params)### Dynamic Hypernetworkdynamic_hypernetwork=JaxDynamicHyperNetwork.from_target(input_dim=32,target_network=target_network,embedding_dim=EMBEDDING_DIM,num_embeddings=NUM_EMBEDDINGS,inputs=jnp.zeros((1,32))# jax needs this to initialize target weights)dynamic_hypernetwork_params=dynamic_hypernetwork.init(key,inp,generate_params_kwargs=dict(x=inp,hidden_state=jnp.zeros((1,32))))# flax needs to initialize hypernetwork parameters first# by default we only output what we'd expect from the target networkoutput=dynamic_hypernetwork.apply(dynamic_hypernetwork_params,inp,generate_params_kwargs=dict(x=inp,hidden_state=jnp.zeros((1,32))))# by default we only output what we'd expect from the target networkoutput=dynamic_hypernetwork.apply(dynamic_hypernetwork_params,inp,generate_params_kwargs=dict(x=inp,hidden_state=jnp.zeros((1,32))))Customizing Hypernetworkshyper-nnmakes it easy to customize and create more complex hypernetworks.The main components to modify are the methodsgenerate_params. This allows for complete control over how the hypernetwork generates parametersFor example, here we extend the linear hypernetwork which uses componentsembedding_moduleandweight_generator. We implement a hypernetwork that could be useful in a multi task setting, where a one hot encoded class embedding is concatenated to every row in the embedding matrix outputted by theembedding_module. In addition, we override both ourmake_embedding_moduleandmake_weight_generatormethods to output customized modules. This whole class implementation is under 50 lines of code!fromtypingimportOptional,Iterable,Any,Tuple,Dictimporttorchimporttorch.nnasnn# static hypernetworkfromhypernn.torchimportTorchHyperNetworkfromhypernn.torch.utilsimportget_weight_chunk_dimsclassMultiTaskHypernetwork(TorchHyperNetwork):def__init__(self,num_tasks:int,target_network:nn.Module,num_target_parameters:Optional[int]=None,embedding_dim:int=100,num_embeddings:int=3,weight_chunk_dim:Optional[int]=None,):super().__init__(target_network=target_network,num_target_parameters=num_target_parameters,)self.num_tasks=num_tasksself.embedding_dim=embedding_dimself.num_embeddings=num_embeddingsself.weight_chunk_dim=weight_chunk_dimifweight_chunk_dimisNone:self.weight_chunk_dim=get_weight_chunk_dims(self.num_target_parameters,num_embeddings)self.embedding_module=self.make_embedding_module()self.weight_generator=self.make_weight_generator()defmake_embedding_module(self)->nn.Module:embedding=nn.Embedding(self.num_embeddings,8)returnnn.Sequential(embedding,nn.Tanh(),nn.Linear(8,self.embedding_dim),nn.Tanh(),)defmake_weight_generator(self)->nn.Module:returnnn.Sequential(nn.Linear(self.embedding_dim+self.num_tasks,32),nn.Tanh(),nn.Linear(32,self.weight_chunk_dim))defgenerate_params(self,one_hot_task_embedding:torch.Tensor)->Tuple[torch.Tensor,Dict[str,Any]]:embedding=self.embedding_module(torch.arange(self.num_embeddings,device=self.device))one_hot_task_embedding=one_hot_task_embedding.repeat(self.num_embeddings,1)# repeat to concat to embeddingconcatenated=torch.cat((embedding,one_hot_task_embedding),dim=-1)generated_params=self.weight_generator(concatenated).view(-1)returngenerated_params,{"embedding":embedding}# usagetarget_network=nn.Sequential(nn.Linear(32,64),nn.ReLU(),nn.Linear(64,32))NUM_TASKS=4EMBEDDING_DIM=4NUM_EMBEDDINGS=32hypernetwork=MultiTaskHypernetwork(num_tasks=NUM_TASKS,target_network=target_network,embedding_dim=EMBEDDING_DIM,num_embeddings=NUM_EMBEDDINGS)inp=torch.zeros((1,32))one_hot_task_embedding=torch.tensor([0.0,0.0,1.0,0.0]).view((1,4))out=hypernetwork(inp,generate_params_kwargs=dict(one_hot_task_embedding=one_hot_task_embedding))Advanced: Using vmap for batching operationsThis is useful when dealing with dynamic hypernetworks that generate different params depending on inputs.Pytorchimporttorch.nnasnnfromfunctorchimportvmap# dynamic hypernetworkfromhypernn.torchimportTorchDynamicHyperNetwork# any moduletarget_network=nn.Sequential(nn.Linear(8,256),nn.ReLU(),nn.Linear(256,32))EMBEDDING_DIM=4NUM_EMBEDDINGS=32# conditioned on input to generate param vectorhypernetwork=TorchDynamicHyperNetwork(target_network=target_network,embedding_dim=EMBEDDING_DIM,num_embeddings=NUM_EMBEDDINGS,input_dim=8)# batch of 10 inputsinp=torch.randn((10,1,8))# use with a for loopoutputs=[]foriinrange(10):outputs.append(hypernetwork(inp[i],generate_params_kwargs=dict(x=inp[i])))outputs=torch.stack(outputs)assertoutputs.size()==(10,1,32)# using vmapfromtypingimportDict,Anydefforward(generated_params,*args,has_aux:bool=False,assert_parameter_shapes:bool=True,generate_params_kwargs:Dict[str,Any]={},**kwargs):returnhypernetwork.forward(*args,generated_params=generated_params,has_aux=has_aux,assert_parameter_shapes=assert_parameter_shapes,generate_params_kwargs=generate_params_kwargs,**kwargs)generated_vmap_params,aux_output=vmap(hypernetwork.generate_params)(inp)outputs=vmap(forward)(generated_vmap_params,inp)assertoutputs.size()==(10,1,32)Future PlansHere's a list of some stuff that will hopefully be added to the library. If anyone has other suggestions, please reach out / create an issue!MNIST exampleLunar Lander ExampleDynamic Hypernetwork ExampleMulti-task Hypernetwork ExampleDedicated documentation websiteEfficient batching for DynamicJaxHypernetworkImplementation ofHyperTransformerImplementation ofRecomposing the Reinforcement Learning Building Blocks with HypernetworksImplementation ofGoal-Conditioned Generators of Deep PoliciesCiting hyper-nnIf you use this software in your publications, please cite it by using the following BibTeX entry.@misc{sudhakaran2022,author={Sudhakaran, Shyam Sudhakaran},title={hyper-nn},howpublished={\url{https://github.com/shyamsn97/hyper-nn}},year={2022},}