Torch-mutable-modules - Use in-place and assignment operations on PyTorch module parameters with support for autograd

Last update: Jun 06, 2022

Related tags

Overview

Torch Mutable Modules

Use 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 of net_2 in-place and backpropagate the gradients to net_1.

# create a parameter predictor network (net_1)
net_1 = torch.nn.Linear(1, 2)

# predict the weights and biases of net_2 using net_1
p_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 values
net_2.weight = p_weight
net_2.bias = p_bias

# compute the output and backpropagate the gradients to net_1
output = 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 the to_mutable_module function.

Installation

You can install torch-mutable-modules from PyPI.

pip install torch-mutable-modules

To upgrade an existing installation of torch-mutable-modules, use the following command:

pip install --upgrade --no-cache-dir torch-mutable-modules

Importing

You can use wildcard imports or import specific functions directly:

# import all functions
from torch_mutable_modules import *

# ... or import the function manually
from torch_mutable_modules import to_mutable_module

Usage

To convert an existing PyTorch module into a mutable module, use the to_mutable_module function:

converted_module = to_mutable_module(
    torch.nn.Linear(1, 1)
) # type of converted_module is still torch.nn.Linear

converted_module.weight *= 0
convreted_module.weight += 69
convreted_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:

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
    
    def forward(self, x):
        return self.linear(x)

my_module = to_mutable_module(MyModule())
my_module.linear.weight *= 0
my_module.linear.weight += 69
my_module.linear.weight # tensor([[69.]], grad_fn=<AddBackward0>)

Usage with CUDA

To create a module on the GPU, simply pass a PyTorch module that is already on the GPU to the to_mutable_module function:

converted_module = to_mutable_module(
    torch.nn.Linear(1, 1).cuda()
) # converted_module is now a mutable module on the GPU

Moving a module to the GPU with .to() and .cuda() after instanciation is NOT supported. Instead, hot-swap the module parameter tensors with their CUDA counterparts.

# both of these are valid
converted_module.weight = converted_module.weight.cuda()
converted_module.bias = converted_module.bias.to("cuda")

Detailed examples

Please check out example.py to see more detailed example usages of the to_mutable_module function.

Contributing

Please feel free to submit issues or pull requests!

You might also like...

A machine learning library for spiking neural networks. Supports training with both torch and jax pipelines, and deployment to neuromorphic hardware.

Rockpool Rockpool is a Python package for developing signal processing applications with spiking neural networks. Rockpool allows you to build network

21 Dec 14, 2022

Implements Stacked-RNN in numpy and torch with manual forward and backward functions

Recurrent Neural Networks Implements simple recurrent network and a stacked recurrent network in numpy and torch respectively. Both flavours implement

1 Nov 16, 2021

A torch.Tensor-like DataFrame library supporting multiple execution runtimes and Arrow as a common memory format

TorchArrow (Warning: Unstable Prototype) This is a prototype library currently under heavy development. It does not currently have stable releases, an

536 Jan 6, 2023

A complete end-to-end demonstration in which we collect training data in Unity and use that data to train a deep neural network to predict the pose of a cube. This model is then deployed in a simulated robotic pick-and-place task.

Object Pose Estimation Demo This tutorial will go through the steps necessary to perform pose estimation with a UR3 robotic arm in Unity. You’ll gain

187 Dec 24, 2022

Torch-mutable-modules - Use in-place and assignment operations on PyTorch module parameters with support for autograd

Related tags

Overview

Torch Mutable Modules

Why does this exist?

Installation

Importing

Usage

Usage with CUDA

Detailed examples

Contributing

You might also like...

A machine learning library for spiking neural networks. Supports training with both torch and jax pipelines, and deployment to neuromorphic hardware.

Implements Stacked-RNN in numpy and torch with manual forward and backward functions

A torch.Tensor-like DataFrame library supporting multiple execution runtimes and Arrow as a common memory format

A complete end-to-end demonstration in which we collect training data in Unity and use that data to train a deep neural network to predict the pose of a cube. This model is then deployed in a simulated robotic pick-and-place task.

Python implementation of MULTIseq barcode alignment using fuzzy string matching and GMM barcode assignment

MM1 and MMC Queue Simulation using python - Results and parameters in excel and csv files

Torch-based tool for quantizing high-dimensional vectors using additive codebooks

Torch implementation of "Enhanced Deep Residual Networks for Single Image Super-Resolution"

Automatic number plate recognition using tech: Yolo, OCR, Scene text detection, scene text recognation, flask, torch

Releases(v1.1.2)

v1.1.2(Jun 6, 2022)

v1.1.1(Jun 5, 2022)

v1.1.0(Feb 1, 2022)

v1.0.1(Feb 1, 2022)

v1.0.0(Feb 1, 2022)

Owner

Kento Nishi

Multi-Anchor Active Domain Adaptation for Semantic Segmentation (ICCV 2021 Oral)

Bayesian Image Reconstruction using Deep Generative Models

UniLM AI - Large-scale Self-supervised Pre-training across Tasks, Languages, and Modalities

Python SDK for building, training, and deploying ML models

Codes for the paper Contrast and Mix: Temporal Contrastive Video Domain Adaptation with Background Mixing

Supervised forecasting of sequential data in Python.

Ensemble Knowledge Guided Sub-network Search and Fine-tuning for Filter Pruning

Joint Learning of 3D Shape Retrieval and Deformation, CVPR 2021

Code of Periodic Activation Functions Induce Stationarity

KeypointDeformer: Unsupervised 3D Keypoint Discovery for Shape Control

CIFAR-10_train-test - training and testing codes for dataset CIFAR-10

Animation of solving the traveling salesman problem to optimality using mixed-integer programming and iteratively eliminating sub tours

It's like Shape Editor in Maya but works with skeletons (transforms).

The code of Zero-shot learning for low-light image enhancement based on dual iteration

Implementation of Neural Distance Embeddings for Biological Sequences (NeuroSEED) in PyTorch

A python implementation of Deep-Image-Analogy based on pytorch.

Face detection using deep learning.

Implementation of E(n)-Transformer, which extends the ideas of Welling's E(n)-Equivariant Graph Neural Network to attention

PointPillars inference with TensorRT

Convert Pytorch model to onnx or tflite, and the converted model can be visualized by Netron