.. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_advanced_numpy_extensions_tutorial.py: Creating Extensions Using numpy and scipy ========================================= **Author**: `Adam Paszke `_ **Updated by**: `Adam Dziedzic` [https://github.com/adam-dziedzic](https://github.com/adam-dziedzic) In this tutorial, we shall go through two tasks: 1. Create a neural network layer with no parameters. - This calls into **numpy** as part of its implementation 2. Create a neural network layer that has learnable weights - This calls into **SciPy** as part of its implementation .. code-block:: python import torch from torch.autograd import Function Parameter-less example ---------------------- This layer doesn’t particularly do anything useful or mathematically correct. It is aptly named BadFFTFunction **Layer Implementation** .. code-block:: python from numpy.fft import rfft2, irfft2 class BadFFTFunction(Function): def forward(self, input): numpy_input = input.detach().numpy() result = abs(rfft2(numpy_input)) return input.new(result) def backward(self, grad_output): numpy_go = grad_output.numpy() result = irfft2(numpy_go) return grad_output.new(result) # since this layer does not have any parameters, we can # simply declare this as a function, rather than as an nn.Module class def incorrect_fft(input): return BadFFTFunction()(input) **Example usage of the created layer:** .. code-block:: python input = torch.randn(8, 8, requires_grad=True) result = incorrect_fft(input) print(result) result.backward(torch.randn(result.size())) print(input) .. rst-class:: sphx-glr-script-out Out: .. code-block:: none tensor([[ 6.4255, 6.8640, 7.4375, 14.2115, 1.1021], [12.8063, 7.0819, 7.5788, 10.9702, 2.2288], [ 5.7878, 2.7235, 11.1733, 8.7793, 13.5364], [10.3233, 11.0394, 9.0303, 2.8435, 13.8431], [ 4.5493, 5.3209, 4.4612, 6.2638, 3.2951], [10.3233, 9.5154, 11.8945, 10.6401, 13.8431], [ 5.7878, 7.6345, 6.3398, 9.3428, 13.5364], [12.8063, 4.7393, 17.3819, 4.3186, 2.2288]], grad_fn=) tensor([[-0.1107, -0.2339, -1.8818, 0.5902, -0.2107, 0.0265, 1.0496, -0.0103], [ 0.1531, 0.4957, -0.3697, -1.2360, -1.6227, -1.5376, 1.3904, -0.6713], [ 0.3373, -1.1663, 1.9184, 1.6435, -0.3209, 1.4998, -0.4512, 0.4015], [ 2.3969, -0.7226, 0.1135, 0.7740, -0.9668, -0.3475, -0.2200, -0.0362], [-0.9314, -0.4943, 0.6745, 1.4668, -1.2265, 0.5994, -0.2864, 2.3817], [ 3.0488, 0.8963, 0.0612, 0.8373, 1.1862, 1.0631, -0.6438, -0.0177], [ 3.3653, -0.5115, -0.3340, -2.1224, 1.1542, 0.0566, -1.1012, -0.2842], [-1.1437, -0.9947, -0.7862, 0.0967, -0.7332, 1.1176, -0.8470, 0.2038]], requires_grad=True) Parametrized example -------------------- In deep learning literature, this layer is confusingly referred to as convolution while the actual operation is cross-correlation (the only difference is that filter is flipped for convolution, which is not the case for cross-correlation). Implementation of a layer with learnable weights, where cross-correlation has a filter (kernel) that represents weights. The backward pass computes the gradient wrt the input and the gradient wrt the filter. .. code-block:: python from numpy import flip import numpy as np from scipy.signal import convolve2d, correlate2d from torch.nn.modules.module import Module from torch.nn.parameter import Parameter class ScipyConv2dFunction(Function): @staticmethod def forward(ctx, input, filter, bias): # detach so we can cast to NumPy input, filter, bias = input.detach(), filter.detach(), bias.detach() result = correlate2d(input.numpy(), filter.numpy(), mode='valid') result += bias.numpy() ctx.save_for_backward(input, filter, bias) return torch.as_tensor(result, dtype=input.dtype) @staticmethod def backward(ctx, grad_output): grad_output = grad_output.detach() input, filter, bias = ctx.saved_tensors grad_output = grad_output.numpy() grad_bias = np.sum(grad_output, keepdims=True) grad_input = convolve2d(grad_output, filter.numpy(), mode='full') # the previous line can be expressed equivalently as: # grad_input = correlate2d(grad_output, flip(flip(filter.numpy(), axis=0), axis=1), mode='full') grad_filter = correlate2d(input.numpy(), grad_output, mode='valid') return torch.from_numpy(grad_input), torch.from_numpy(grad_filter).to(torch.float), torch.from_numpy(grad_bias).to(torch.float) class ScipyConv2d(Module): def __init__(self, filter_width, filter_height): super(ScipyConv2d, self).__init__() self.filter = Parameter(torch.randn(filter_width, filter_height)) self.bias = Parameter(torch.randn(1, 1)) def forward(self, input): return ScipyConv2dFunction.apply(input, self.filter, self.bias) **Example usage:** .. code-block:: python module = ScipyConv2d(3, 3) print("Filter and bias: ", list(module.parameters())) input = torch.randn(10, 10, requires_grad=True) output = module(input) print("Output from the convolution: ", output) output.backward(torch.randn(8, 8)) print("Gradient for the input map: ", input.grad) .. rst-class:: sphx-glr-script-out Out: .. code-block:: none Filter and bias: [Parameter containing: tensor([[-0.7778, 2.2590, -0.9142], [ 0.9111, -1.3104, -0.6073], [-0.3173, -0.1487, -0.3237]], requires_grad=True), Parameter containing: tensor([[-0.5367]], requires_grad=True)] Output from the convolution: tensor([[-0.7660, -2.4804, -1.7901, 2.0918, -4.0412, 1.6315, 0.8086, 0.8650], [-0.3624, -0.5671, -0.1210, -0.2485, 2.1909, -1.3427, -3.4609, -2.3622], [ 1.1806, 1.1542, 3.1647, -7.2220, 0.3783, -5.2158, 1.8066, 2.1748], [ 1.5890, -2.9704, -2.1720, 2.0465, -1.0344, 2.3148, -2.6923, 1.4638], [-3.3533, 0.6709, 0.1915, -1.0633, 1.2107, -5.1670, 1.2944, 0.2674], [-0.8958, 1.2640, 0.4206, -1.9615, -3.1485, 0.8377, 5.4179, -0.4798], [ 3.9596, 2.6956, -4.5437, 4.9535, 0.4918, -1.5233, -6.6145, -0.3701], [ 4.9656, -6.5131, 6.2981, -2.3091, 0.2646, -3.4335, 2.1812, -0.7152]], grad_fn=) Gradient for the input map: tensor([[-0.3006, 1.7584, -4.5034, 4.5046, 1.8851, -3.0384, 0.6406, 1.0668, 1.9149, -1.1734], [ 0.0706, -1.9065, 6.7826, -3.2842, -4.0473, 1.0767, -0.7694, -1.7111, -1.7290, -0.6216], [-0.5300, 3.5590, -5.0111, 0.5257, -3.0366, 1.9396, -0.2996, -0.7462, 1.0946, -0.1694], [-0.6314, 3.6796, -6.5135, -1.0524, 5.1298, -0.6599, 2.1989, 0.9593, 2.8613, -0.5312], [ 2.5064, -8.2649, 3.4903, 3.1351, 0.9233, -0.4950, -2.4464, 0.4043, -1.9103, 0.2270], [-3.0514, 4.6734, 3.9245, -2.1171, -3.6524, -0.8866, 0.2607, -0.4007, -0.2227, 0.6759], [ 0.8078, 3.0731, -6.0957, -1.3559, 2.7699, 0.2014, 2.8031, -1.9639, -3.0835, 2.6467], [ 0.7234, -3.0810, 0.5078, -1.7652, 1.0972, 1.7048, -3.5835, -3.9572, 4.2907, 2.0443], [-0.4453, 0.1479, -1.2103, 1.3386, -0.0397, -2.6211, -1.1376, 2.8240, 2.4993, 1.0581], [ 0.0188, -0.0257, 0.2433, -0.1165, -0.0818, -0.0987, 0.4207, 0.6437, 0.6522, 0.1691]]) **Check the gradients:** .. code-block:: python from torch.autograd.gradcheck import gradcheck moduleConv = ScipyConv2d(3, 3) input = [torch.randn(20, 20, dtype=torch.double, requires_grad=True)] test = gradcheck(moduleConv, input, eps=1e-6, atol=1e-4) print("Are the gradients correct: ", test) .. rst-class:: sphx-glr-script-out Out: .. code-block:: none Are the gradients correct: True **Total running time of the script:** ( 0 minutes 0.319 seconds) .. _sphx_glr_download_advanced_numpy_extensions_tutorial.py: .. only :: html .. container:: sphx-glr-footer :class: sphx-glr-footer-example .. container:: sphx-glr-download :download:`Download Python source code: numpy_extensions_tutorial.py ` .. container:: sphx-glr-download :download:`Download Jupyter notebook: numpy_extensions_tutorial.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_