PyTorch 使用 numpy 和 scipy 创建扩展

原文： https://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html

作者： Adam Paszke

已由更新： Adam Dziedzic

在本教程中，我们将完成两个任务：

1. 创建没有参数的神经网络层。> 调用 *numpy 作为其实现的一部分
2. 创建具有可学习权重的神经网络层> 作为其实现的一部分，它引发 *SciPy

import torch
from torch.autograd import Function

无参数的例子

该层并没有做任何有用的或数学上正确的事情。

它被恰当地命名为 BadFFTFunction

层实现：

from numpy.fft import rfft2, irfft2


class BadFFTFunction(Function):
    @staticmethod
    def forward(ctx, input):
        numpy_input = input.detach().numpy()
        result = abs(rfft2(numpy_input))
        return input.new(result)


    @staticmethod
    def backward(ctx, 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.apply(input)

创建的图层的用法示例：

input = torch.randn(8, 8, requires_grad=True)
result = incorrect_fft(input)
print(result)
result.backward(torch.randn(result.size()))
print(input)

得出：

tensor([[ 3.4193,  6.2492,  1.4814,  3.9813, 14.7739],
        [ 5.8865,  2.3533,  1.4765,  4.1547,  5.1887],
        [ 6.9061,  4.7597,  5.9473, 11.8609, 10.0465],
        [10.6724,  6.5964, 11.6951,  3.7400,  9.3114],
        [ 4.9715,  5.7563, 10.5443,  4.6442,  2.3052],
        [10.6724, 12.7221,  1.7242,  7.2647,  9.3114],
        [ 6.9061,  8.1754,  4.0003,  3.1523, 10.0465],
        [ 5.8865, 15.8094,  6.9266,  3.8533,  5.1887]],
       grad_fn=<BadFFTFunctionBackward>)
tensor([[ 0.4963,  0.5220, -0.0559,  0.2241,  2.1238,  0.0324, -0.1219,  0.7318],
        [ 1.0476, -0.1841, -0.2717,  0.5470, -1.3174, -0.6817,  1.0102, -0.2014],
        [-1.1131,  1.0054,  0.3593,  0.6158,  0.5398,  2.3020, -0.1305, -0.8611],
        [-0.4693, -1.3720, -0.8196, -1.4975, -0.4474, -0.1150, -0.3285,  1.8079],
        [-0.2928,  1.4333, -0.2744, -0.9194, -0.2592,  0.4996,  0.7862,  0.3972],
        [ 0.7595,  0.5625, -0.7585, -0.1439, -0.5243, -1.0789,  0.4915,  1.5880],
        [-0.6971,  0.0267,  0.2316, -0.8939, -1.9865, -0.7424, -0.6252,  0.8415],
        [-0.9989,  1.0916,  0.2223,  2.1130, -0.3831,  0.9612, -1.8703,  0.4848]],
       requires_grad=True)

参数化示例

在深度学习文献中，该层被混淆地称为卷积，而实际操作是互相关的(唯一的区别是滤波器被卷积以进行卷积，而互相关不是这种情况）。

具有可学习的权重的层的实现，其中互相关具有表示权重的过滤器(内核）。

反向通过计算输入的梯度和过滤器的梯度。

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)

用法示例：

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)

得出:

Filter and bias:  [Parameter containing:
tensor([[-0.4935, -1.9784, -0.7520],
        [ 0.0575, -1.3029, -1.9318],
        [ 1.1692,  0.3187, -0.3044]], requires_grad=True), Parameter containing:
tensor([[-0.9122]], requires_grad=True)]
Output from the convolution:  tensor([[ 1.9228e+00,  5.1708e+00, -1.7722e-01,  6.0846e-01,  3.1403e+00,
          2.8270e+00, -2.7449e+00, -7.5212e+00],
        [ 9.2997e-01,  1.4772e-01, -1.7859e+00, -4.8096e+00, -3.2894e+00,
          2.9466e+00,  2.2201e+00, -1.7857e+00],
        [-3.4693e+00,  1.0293e-01, -1.6728e+00, -6.0165e+00, -6.5472e+00,
         -8.0421e-01, -2.2059e+00, -7.3037e+00],
        [-7.1793e+00,  2.5055e-01,  1.0368e+00, -1.6319e+00, -8.3195e+00,
         -4.6471e+00, -1.6651e+00, -3.5958e+00],
        [-1.3189e+00,  1.9218e+00,  4.4448e+00,  2.2613e+00, -4.9712e+00,
         -7.1531e+00, -7.3305e-01, -1.7002e+00],
        [ 8.7136e-01,  6.6128e-01,  2.9627e+00,  1.1253e+00, -4.1352e+00,
         -2.5405e+00,  2.5132e+00,  5.8881e+00],
        [ 1.4052e+00, -2.3034e+00, -8.8434e-01, -5.1262e-02, -1.0159e+01,
         -5.9782e+00, -2.7782e-01,  1.6121e+00],
        [ 2.0402e+00,  1.7031e-03,  1.5932e+00,  4.6277e+00, -8.4549e+00,
         -1.4017e+01, -6.9038e+00, -2.8451e+00]],
       grad_fn=<ScipyConv2dFunctionBackward>)
Gradient for the input map:  tensor([[ 3.5673e-01,  1.5783e+00,  5.5266e-01, -2.2223e+00, -1.1648e+00,
          1.3576e-01, -9.2908e-01, -9.5302e-01, -1.0746e-01,  4.7559e-02],
        [-8.7051e-01, -2.4345e+00,  2.0736e-01, -1.8105e+00, -2.6859e+00,
         -1.2512e+00, -3.8021e-01, -6.9782e-01,  2.8626e+00,  1.4599e+00],
        [-1.1194e+00, -4.4828e+00, -2.9077e+00,  1.2989e+00,  7.7022e-01,
          1.9482e+00,  8.1654e-01,  4.3522e+00,  4.5537e+00,  3.8784e+00],
        [ 2.0067e+00,  3.9783e-01,  5.8589e-01, -1.2962e+00, -1.7774e+00,
          2.7983e-01,  2.0063e+00, -5.3575e-01,  1.9437e+00,  9.2401e-01],
        [ 9.6117e-01,  1.1666e+00,  2.1162e+00,  2.9779e+00, -4.0048e+00,
         -2.1271e-01,  3.6012e-01, -1.5548e-03, -2.2730e+00, -1.9148e+00],
        [-5.5450e-01, -3.5749e+00,  2.1108e-01,  2.7341e+00, -5.4415e-01,
         -1.0326e-01,  5.2842e+00,  6.0280e+00,  1.7129e+00, -1.0743e+00],
        [-4.5225e-01, -3.2305e+00, -3.9962e+00, -1.4418e+00, -6.6439e+00,
         -6.6198e+00, -4.7138e+00,  2.4043e-01,  3.5485e+00, -2.6767e-01],
        [ 1.2502e+00,  1.5484e-01,  4.7037e+00,  4.7502e+00, -2.7102e-01,
         -5.2364e+00, -7.6462e+00, -1.5124e+00,  3.1666e+00,  1.0458e+00],
        [ 7.6941e-01,  1.0242e-01,  3.4887e+00,  6.9709e+00,  4.7607e+00,
          1.1723e+00, -2.5379e+00,  1.0167e+00,  3.5257e+00,  2.5870e+00],
        [ 1.8773e-01, -2.9968e+00, -2.2132e+00, -2.2993e-01,  2.0811e+00,
          5.9679e-01, -1.6988e+00, -1.8538e+00, -1.1920e-01,  4.0583e-01]])

检查渐变：

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)

得出:

Are the gradients correct:  True

脚本的总运行时间：(0 分钟 4.206 秒）

Download Python source code: numpy_extensions_tutorial.py Download Jupyter notebook: numpy_extensions_tutorial.ipynb