# Implementing a custom convolution using conv2d\_input and conv2d\_weight

**URL:** <https://discuss.pytorch.org/t/implementing-a-custom-convolution-using-conv2d-input-and-conv2d-weight/18556>\
**Category:** Uncategorized\
**Created:** [May 23, 2018, 8:27am UTC](https://discuss.pytorch.org/t/implementing-a-custom-convolution-using-conv2d-input-and-conv2d-weight/18556 "2018-05-23T08:27:36Z")\
**Posts on this page:** 1\
**Showing post:** 24

<div class="post-metadata">

**Author:** ![arman-yekkehkhani](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/arman-yekkehkhani/32/32661_2.png) [@arman-yekkehkhani](https://discuss.pytorch.org/u/arman-yekkehkhani)\
**Post date:** [July 8, 2021, 7:43pm UTC](https://discuss.pytorch.org/t/implementing-a-custom-convolution-using-conv2d-input-and-conv2d-weight/18556/24 "2021-07-08T19:43:56Z")

</div>

Thanks to all the other responses. Here is the memory-efficient, fast solution:

```auto
from torch.utils.cpp_extension import load

cudnn_convolution = load(name="cudnn_convolution", sources=["cudnn_convolution.cpp"], verbose=True)

class CustomConv2d(Function):
    @staticmethod
    def forward(ctx, input, weight, bias, stride, padding, dilation, groups):
        ctx.save_for_backward(input, weight, bias)
        ctx.conf = {
            "stride": stride,
            "padding": padding,
            "dilation": dilation,
            "groups": groups
        }
        return cudnn_convolution.convolution(input, weight, bias, stride, padding, dilation, groups,
                                             False, False)

    @staticmethod
    def backward(ctx, grad_output):
        input, weight, bias = ctx.saved_variables
        conf = ctx.conf
        input_grad = weight_grad = bias_grad = stride_grad = padding_grad = dilation_grad = groups_grad = None
        if ctx.needs_input_grad[0]:
            input_grad = cudnn_convolution.convolution_backward_input(input.shape, weight, grad_output, conf["stride"],
                                                                      conf["padding"], conf["dilation"], conf["groups"],
                                                                      False, False, False)
        if ctx.needs_input_grad[1]:
            weight_grad = cudnn_convolution.convolution_backward_weight(input, weight.shape, grad_output,
                                                                        conf["stride"], conf["padding"],
                                                                        conf["dilation"], conf["groups"],
                                                                        False, False, False)

        if bias is not None and ctx.needs_input_grad[2]:
            bias_grad = grad_output.sum(dim=(0, 2, 3))

        return input_grad, weight_grad, bias_grad, stride_grad, padding_grad, dilation_grad, groups_grad

```

For the above code to work, you need to create a file named “cudnn\_convolution.cpp” which containts the following code(copied from this repo [PyTorch cuDNN Convolution](https://github.com/jordan-g/PyTorch-cuDNN-Convolution)):

```auto
#include <torch/extension.h>
#include <vector>
#include <ATen/NativeFunctions.h>
#include <ATen/Config.h>

/*
PyTorch extension enabling direct access to the following cuDNN-accelerated C++ functions
that are included in PyTorch:
    - cudnn_convolution
    - cudnn_convolution_backward_weight
    - cudnn_convolution_backward_input
The functions defined here can be called from Python in replacement of
torch.nn.conv2d, torch.nn.grad.conv2d_weight and torch.nn.grad.conv2d_input,
and run significantly faster. See 'example.py' for how these functions
are called.
Adapted from code posted by hanspinckaers:
https://discuss.pytorch.org/t/cuda-error-with-cudnn-convolution-backward-weight-function/41214
*/

at::Tensor convolution(
    const at::Tensor& input,
    const at::Tensor& weight,
    const at::Tensor& bias,
    c10::ArrayRef<int64_t> stride,
    c10::ArrayRef<int64_t> padding,
    c10::ArrayRef<int64_t> dilation,
    int64_t groups,
    bool benchmark,
    bool deterministic) {

    return at::cudnn_convolution(
        input,
        weight,
        bias,
        padding,
        stride,
        dilation,
        groups,
        benchmark,
        deterministic); 
}

at::Tensor convolution_backward_weight(
    const at::Tensor& input,
    c10::ArrayRef<int64_t> weight_size,
    const at::Tensor& grad_output,
    c10::ArrayRef<int64_t> stride,
    c10::ArrayRef<int64_t> padding,
    c10::ArrayRef<int64_t> dilation,
    int64_t groups,
    bool benchmark,
    bool deterministic,
    bool allow_tf32) {

    return at::cudnn_convolution_backward_weight(
        weight_size,
        grad_output,
        input,
        padding,
        stride,
        dilation,
        groups,
        benchmark,
        deterministic,
        allow_tf32);
}

at::Tensor convolution_backward_input(
    c10::ArrayRef<int64_t> input_size,
    const at::Tensor& weight,
    const at::Tensor& grad_output,
    c10::ArrayRef<int64_t> stride,
    c10::ArrayRef<int64_t> padding,
    c10::ArrayRef<int64_t> dilation,
    int64_t groups,
    bool benchmark,
    bool deterministic,
    bool allow_tf32) {

    return at::cudnn_convolution_backward_input(
        input_size,
        grad_output,
        weight,
        padding,
        stride,
        dilation,
        groups,
        benchmark,
        deterministic,
        allow_tf32);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("convolution", &convolution, "convolution");
    m.def("convolution_backward_weight", &convolution_backward_weight, "convolution backward weight");
    m.def("convolution_backward_input", &convolution_backward_input, "convolution backward input");
}

```

---

_[View the full topic](https://discuss.pytorch.org/t/implementing-a-custom-convolution-using-conv2d-input-and-conv2d-weight/18556)._