Package: torch-mutable-modules
Torch Mutable ModulesUse in-place and assignment operations on PyTorch module parameters with support for autograd.Why does this exist?PyTorch does not allow in-place operations on module parameters (usually desirable):linear_layer=torch.nn.Linear(1,1)linear_layer.weight.data+=69# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^# Valid, but will NOT store grad_fn=<AddBackward0>linear_layer.weight+=420# ^^^^^^^^^^^^^^^^^^^^^^^^# RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.In some cases, however, it is useful to be able to modify module parameters in-place. For example, if we have a neural network (net_1) that predicts the parameter values to another neural network (net_2), we need to be able to modify the weights ofnet_2in-place and backpropagate the gradients tonet_1.# create a parameter predictor network (net_1)net_1=torch.nn.Linear(1,2)# create an optimizer for net_1optimizer_1=torch.optim.SGD(net_1.parameters(),lr=0.01)# predict the weights and biases of net_2 using net_1p_weight_and_bias=net_1(input_0).unsqueeze(2)p_weight,p_bias=p_weight_and_bias[:,0],p_weight_and_bias[:,1]# create a mutable network (net_2)net_2=to_mutable_module(torch.nn.Linear(1,1))# hot-swap the weights and biases of net_2 with the predicted valuesnet_2.weight=p_weightnet_2.bias=p_bias# compute the output and backpropagate the gradients to net_1output=net_2(input_1)loss=criterion(output,label)loss.backward()optimizer.step()This library provides a way to easily convert PyTorch modules into mutable modules with theto_mutable_modulefunction.InstallationYou can installtorch-mutable-modulesfromPyPI.pipinstalltorch-mutable-modulesTo upgrade an existing installation oftorch-mutable-modules, use the following command:pipinstall--upgrade--no-cache-dirtorch-mutable-modulesImportingYou can use wildcard imports or import specific functions directly:# import all functionsfromtorch_mutable_modulesimport*# ... or import the function manuallyfromtorch_mutable_modulesimportto_mutable_moduleUsageTo convert an existing PyTorch module into a mutable module, use theto_mutable_modulefunction:converted_module=to_mutable_module(torch.nn.Linear(1,1))# type of converted_module is still torch.nn.Linearconverted_module.weight*=0convreted_module.weight+=69convreted_module.weight# tensor([[69.]], grad_fn=<AddBackward0>)You can also declare your own PyTorch module classes as mutable, and all child modules will be recursively converted into mutable modules:classMyModule(nn.Module):def__init__(self):super().__init__()self.linear=nn.Linear(1,1)defforward(self,x):returnself.linear(x)my_module=to_mutable_module(MyModule())my_module.linear.weight*=0my_module.linear.weight+=69my_module.linear.weight# tensor([[69.]], grad_fn=<AddBackward0>)Usage with CUDATo create a module on the GPU, simply pass a PyTorch module that is already on the GPU to theto_mutable_modulefunction:converted_module=to_mutable_module(torch.nn.Linear(1,1).cuda())# converted_module is now a mutable module on the GPUMoving a module to the GPU with.to()and.cuda()after instanciation isNOTsupported. Instead, hot-swap the module parameter tensors with their CUDA counterparts.# both of these are validconverted_module.weight=converted_module.weight.cuda()converted_module.bias=converted_module.bias.to("cuda")Detailed examplesPlease check outexample.pyto see more detailed example usages of theto_mutable_modulefunction.ContributingPlease feel free to submit issues or pull requests!