使用 numpy 和 scipy 创建扩展

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

注意

单击此处的下载完整的示例代码

作者： Adam Paszke

已由更新： Adam Dziedzic

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

  1. 创建没有参数的神经网络层。

    > 调用 *numpy 作为其实现的一部分

  2. 创建具有可学习权重的神经网络层

    > 作为其实现的一部分，它引发 *SciPy

  1. import torch
  2. from torch.autograd import Function

无参数的例子

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

它被恰当地命名为 BadFFTFunction

层实现

  1. from numpy.fft import rfft2, irfft2
  3. class BadFFTFunction(Function):
  4.     @staticmethod
  5.     def forward(ctx, input):
  6.         numpy_input = input.detach().numpy()
  7.         result = abs(rfft2(numpy_input))
  8.         return input.new(result)
  10.     @staticmethod
  11.     def backward(ctx, grad_output):
  12.         numpy_go = grad_output.numpy()
  13.         result = irfft2(numpy_go)
  14.         return grad_output.new(result)
  16. # since this layer does not have any parameters, we can
  17. # simply declare this as a function, rather than as an nn.Module class
  19. def incorrect_fft(input):
  20.     return BadFFTFunction.apply(input)

创建的图层的用法示例：

  1. input = torch.randn(8, 8, requires_grad=True)
  2. result = incorrect_fft(input)
  3. print(result)
  4. result.backward(torch.randn(result.size()))
  5. print(input)

出：

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

参数化示例

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

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

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

  1. from numpy import flip
  2. import numpy as np
  3. from scipy.signal import convolve2d, correlate2d
  4. from torch.nn.modules.module import Module
  5. from torch.nn.parameter import Parameter
  7. class ScipyConv2dFunction(Function):
  8.     @staticmethod
  9.     def forward(ctx, input, filter, bias):
  10.         # detach so we can cast to NumPy
  11.         input, filter, bias = input.detach(), filter.detach(), bias.detach()
  12.         result = correlate2d(input.numpy(), filter.numpy(), mode='valid')
  13.         result += bias.numpy()
  14.         ctx.save_for_backward(input, filter, bias)
  15.         return torch.as_tensor(result, dtype=input.dtype)
  17.     @staticmethod
  18.     def backward(ctx, grad_output):
  19.         grad_output = grad_output.detach()
  20.         input, filter, bias = ctx.saved_tensors
  21.         grad_output = grad_output.numpy()
  22.         grad_bias = np.sum(grad_output, keepdims=True)
  23.         grad_input = convolve2d(grad_output, filter.numpy(), mode='full')
  24.         # the previous line can be expressed equivalently as:
  25.         # grad_input = correlate2d(grad_output, flip(flip(filter.numpy(), axis=0), axis=1), mode='full')
  26.         grad_filter = correlate2d(input.numpy(), grad_output, mode='valid')
  27.         return torch.from_numpy(grad_input), torch.from_numpy(grad_filter).to(torch.float), torch.from_numpy(grad_bias).to(torch.float)
  29. class ScipyConv2d(Module):
  30.     def __init__(self, filter_width, filter_height):
  31.         super(ScipyConv2d, self).__init__()
  32.         self.filter = Parameter(torch.randn(filter_width, filter_height))
  33.         self.bias = Parameter(torch.randn(1, 1))
  35.     def forward(self, input):
  36.         return ScipyConv2dFunction.apply(input, self.filter, self.bias)

用法示例：

  1. module = ScipyConv2d(3, 3)
  2. print("Filter and bias: ", list(module.parameters()))
  3. input = torch.randn(10, 10, requires_grad=True)
  4. output = module(input)
  5. print("Output from the convolution: ", output)
  6. output.backward(torch.randn(8, 8))
  7. print("Gradient for the input map: ", input.grad)

Out:

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

检查渐变：

  1. from torch.autograd.gradcheck import gradcheck
  3. moduleConv = ScipyConv2d(3, 3)
  5. input = [torch.randn(20, 20, dtype=torch.double, requires_grad=True)]
  6. test = gradcheck(moduleConv, input, eps=1e-6, atol=1e-4)
  7. print("Are the gradients correct: ", test)

Out:

  1. Are the gradients correct:  True

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

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

由狮身人面像画廊生成的画廊