Torch Pconv

Faster and more memory efficient implementation of the Partial Convolution 2D layer in PyTorch equivalent to the standard Nvidia implementation.

This implementation has numerous advantages:

1. It is strictly equivalent in computation to the reference implementation by Nvidia . I made unit tests to assess that all throughout development.
2. It's commented and more readable
3. It's faster and more memory efficient, which means you can use more layers on smaller GPUs. It's a good thing considering today's GPU prices.
4. It's a PyPI-published library. You can pip install it instead of copy/pasting source code, and get the benefit of ( free) bugfixes when someone notice a bug in the implementation.

Getting started

pip3 install torch_pconv

Usage

import torch
from torch_pconv import PConv2d

images = torch.rand(32, 3, 256, 256)
masks = (torch.rand(32, 256, 256) > 0.5).to(torch.float32)

pconv = PConv2d(
    in_channels=3,
    out_channels=64,
    kernel_size=7,
    stride=1,
    padding=2,
    dilation=2,
    bias=True
)

output, shrunk_masks = pconv(images, masks)

Performance improvement

Test

You can find the reference implementation by Nvidia here .

I tested their implementation vs mine one the following configuration:

The goal here was to produce the most computationally expensive partial convolution operator so that the performance difference is displayed better.

I compute both the forward and the backward pass, in case one consumes more memory than the other.

Results

torch_pconv

Development

To install the latest version from Github, run:

git clone git@github.com:DesignStripe/torch_pconv.git torch_pconv
cd torch_pconv
pip3 install -U .