PyTorch Custom Convolution Layer

This project is about how to define a custom convolution layer in PyTorch, and use CUDA function to implement convolution.

Content

/cpp C++ extension include CUDA interface and Python module bind.

/cuda Implicit gemm convolution implementation.

/include Declaration about forward/backward convolution.

/pytorch Include setup.py script, custom convolution layer definition.

Build

$ sh setup.sh

If you don't have root permission, add environment option --prefix="/home/user/.conda/envs/yourenvname/.

Run

from conv_layer import Conv2d

Conv2d(in_channels = 1,out_channels = 16,kernel_size = 3,stride = 1,padding = 1)

To use custom conv layer, just import conv_layer.Conv2d, and use it like nn.Conv2d.