Faster and more memory efficient implementation of the Partial Convolution 2D layer in PyTorch equivalent to the standard NVidia implem.
Project description
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:
- It is strictly equivalent in computation to the reference implementation by Nvidia . I made unit tests to assess that all throughout development.
- It's commented and more readable
- 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.
- 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:
Parameter | Value |
---|---|
in_channels | 64 |
out_channels | 128 |
kernel_size | 9 |
stride | 1 |
padding | 3 |
bias | True |
input height/width | 256 |
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 | Nvidia® (Guilin) | |
---|---|---|
Forward only | 813 466 624 | 4 228 120 576 |
Backward only | 1 588 201 480 | 1 588 201 480 |
Forward + Backward | 2 405 797 640 | 6 084 757 512 |
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 .
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for torch_pconv-0.1.0-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 01adcaca3e432c09a4702c8212f72eebca0a9fed4c5d42000cd77dd46a67e678 |
|
MD5 | dde7c740abd8f84ae862bbf18c3381b0 |
|
BLAKE2b-256 | a1347106ceba8be52f0b749a8adc0d30d006f4669e9ef4e5a4eed75fad8a7fed |