使用 NumPy 和 SciPy 创建扩展#

创建于：2017年3月24日 | 最后更新于：2023年4月25日 | 最后验证：未验证

作者： Adam Paszke

更新者： Adam Dziedzic

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

创建一个没有参数的神经网络层。
- 这在其实现中调用了 numpy
创建一个具有可学习权重的神经网络层
- 这在其实现中调用了 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([[ 6.3896,  3.9308, 10.6650, 15.2876,  7.4706],
        [ 4.8754, 13.0552, 10.6457,  4.4485,  9.3017],
        [ 8.8655, 15.9728,  3.6280, 11.2058,  5.2940],
        [ 6.4814,  8.1925,  8.1897,  1.7061,  5.2642],
        [ 1.4810,  5.7961,  6.1371, 14.4215,  2.5298],
        [ 6.4814,  3.3904,  5.8236, 20.3840,  5.2642],
        [ 8.8655,  8.2430,  4.5068, 10.9887,  5.2940],
        [ 4.8754,  7.8300,  0.7459, 10.5094,  9.3017]],
       grad_fn=<BadFFTFunctionBackward>)
tensor([[ 0.6020,  3.0631, -0.3781,  1.0669, -0.3826, -1.4503,  1.0378, -0.2777],
        [ 1.4850,  1.0758,  0.9505, -1.5323, -0.4903,  0.1031,  0.4727,  0.4069],
        [-0.2958,  1.5686, -0.8486, -0.0322,  0.0599, -0.0990,  2.3614, -0.3482],
        [-0.5261,  1.6887, -1.6335, -0.8338,  0.4496, -1.1807,  1.0672, -0.3636],
        [-0.2068,  0.5354,  0.1173, -1.0162,  1.0584,  0.2440, -1.6391, -1.0929],
        [ 1.3857, -0.3262,  0.3405,  1.2999, -1.4753,  2.1700,  0.1948, -0.4542],
        [-1.2610,  0.4796, -2.6763, -0.7589,  2.4776,  1.0007,  0.7064,  0.3197],
        [-0.5643,  0.0874, -0.5357,  2.1663, -1.2133, -0.3423, -1.1803, -0.2375]],
       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.9934, -1.8313,  1.1618],
        [-0.6469,  1.5618,  1.0743],
        [-0.6961, -0.2246,  0.7994]], requires_grad=True), Parameter containing:
tensor([[0.5140]], requires_grad=True)]
Output from the convolution:  tensor([[ 0.7987, -1.8356, -0.0368,  4.0213,  2.4479,  0.0126,  0.4510, -4.8292],
        [ 0.3330,  3.3282, -4.0302, -3.8571, -0.6718,  1.1097, -1.7532,  6.8066],
        [ 2.4253, -0.7357, -0.9237,  1.1130,  2.5390,  1.2486, -4.5746, -2.3091],
        [ 5.3401, -5.0762,  3.5291, -1.5899,  4.9255,  3.3578,  0.5538,  6.7068],
        [-5.6698, -0.5987, -4.5077, -2.6584,  1.0617, -1.1368, -2.0391, -3.0442],
        [-0.9867,  4.8055,  5.7074, -0.3322,  5.3515, -0.3227, -0.4573,  0.1265],
        [-5.3482,  2.1900, -6.0843, -0.8190,  1.3040,  8.3411,  0.1125, -6.6049],
        [-3.4074,  2.8798,  4.8524, -3.1583,  6.4120, -6.4762,  1.5767,  7.5515]],
       grad_fn=<ScipyConv2dFunctionBackward>)
Gradient for the input map:  tensor([[ 1.9307,  2.1078, -4.6170,  0.8990, -2.4093,  3.6706,  2.9916,  1.3961,
         -2.0263,  0.1481],
        [-0.6371, -7.6767,  1.6163, -2.1760,  1.3180,  6.4428, -0.7031, -6.2453,
         -0.2714, -0.1368],
        [-0.0892,  1.5641, -1.0731, -0.8784,  4.2021,  0.1815, -4.9030,  0.6106,
          2.4680, -1.8023],
        [-0.6922,  1.6504,  0.4926, -1.9421,  3.0568,  6.2646,  0.1321,  0.8156,
          0.2104, -2.0462],
        [-0.5986, -2.1367, -2.4757, -3.0877,  4.8412, -4.4043, -4.5396,  1.0639,
          0.7828, -0.7707],
        [ 0.3211,  1.9110, -4.3512,  4.3929,  2.7580,  4.2731,  0.9912, -0.1901,
          2.5576,  0.2017],
        [-0.0427,  1.1144,  6.3763, -0.1272, -0.3100, -0.3756, -5.2695,  0.7732,
          4.3270, -2.0504],
        [ 0.6700,  0.6825, -2.1998, -5.0459,  2.1912, -2.3359, -4.2094,  2.9008,
         -5.2431, -1.0019],
        [ 0.3534,  1.5500, -0.9943,  0.4659, -0.1893, -2.5566,  1.4202,  1.9999,
          1.9748, -0.4146],
        [ 0.1057,  0.5506, -0.7871, -0.3300,  1.4084, -1.6683, -0.3937,  0.6428,
         -0.5933,  0.8872]])

检查梯度

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 分 0.599 秒)

使用 NumPy 和 SciPy 创建扩展#

无参数示例#

带参数示例#

文档

教程

资源