使用 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([[ 1.7967, 11.4369,  6.1887,  5.7943,  1.6766],
        [ 1.4158, 17.9697,  8.1843,  6.8669,  4.8112],
        [ 8.5268,  8.6887, 10.4486,  3.7192,  0.5529],
        [ 4.5089,  4.6417, 10.4153,  4.8813,  5.1072],
        [11.3335, 10.8583,  2.0834, 15.1015,  1.0625],
        [ 4.5089,  6.2905,  4.6241, 11.5325,  5.1072],
        [ 8.5268, 12.4643,  9.4872,  5.5986,  0.5529],
        [ 1.4158,  6.9435,  4.0634,  5.1742,  4.8112]],
       grad_fn=<BadFFTFunctionBackward>)
tensor([[-0.4744,  0.8724, -0.6717, -0.6049,  0.7731, -1.2441, -1.3551, -0.0614],
        [-1.0561, -1.2702,  0.8774,  2.7034,  0.6695,  1.6789, -0.0487, -0.7222],
        [-1.0487, -1.3456,  1.7747,  0.9599, -0.0102,  1.6475,  0.3243, -0.8861],
        [-0.0168,  0.3763, -0.3420,  1.5060,  0.4426, -0.1962, -0.8757, -1.5912],
        [ 1.1167, -0.3605,  0.1641, -1.5479,  0.0084, -1.4061, -1.5425,  1.5558],
        [ 1.7365,  0.3380,  0.7479,  0.9748,  0.5120,  0.0205, -0.0784, -0.2729],
        [-1.5305, -0.7548, -0.1640, -0.9805,  0.3473, -0.2693,  0.5891,  1.3566],
        [-0.8834, -1.5920,  1.0963,  0.5275, -0.6533,  0.5858,  1.3082,  0.0626]],
       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.1385, -0.5306,  0.0377],
        [ 1.0868, -1.6797,  0.9640],
        [ 1.2143, -0.1341, -0.8983]], requires_grad=True), Parameter containing:
tensor([[-1.4071]], requires_grad=True)]
Output from the convolution:  tensor([[-1.9088,  0.8294, -1.7339, -1.3235, -0.6347, -3.5310, -0.7351, -4.2402],
        [-2.1465, -0.2379, -0.8618, -5.1499,  0.3606,  1.0557, -1.9359, -4.9706],
        [-3.3424, -3.9605,  5.0766, -1.1647, -4.6563, -0.2947,  1.4976, -8.4955],
        [-2.6366, -8.1105, -1.1617, -0.7926,  0.3108, -1.2655,  1.9872, -5.8606],
        [ 0.8425, -2.2686,  0.1502,  0.7788, -0.5958, -5.7892,  2.6934, -1.4429],
        [-0.8635, -4.5383, -0.0142, -0.3470, -2.2123, -0.3513, -6.9511,  1.4569],
        [ 0.1413, -2.9592, -1.3666,  1.5642, -8.7162,  4.0966, -2.6379, -1.3338],
        [-3.5865, -1.5555, -2.1811,  1.2695, -3.6870,  0.2894, -0.3669, -3.2512]],
       grad_fn=<ScipyConv2dFunctionBackward>)
Gradient for the input map:  tensor([[ 0.0593, -0.3022,  0.0403,  0.8574,  0.6406, -0.8675,  0.0582,  0.5035,
         -1.0287,  0.0718],
        [ 0.7005, -2.3835, -0.1653,  2.4221,  0.4658, -3.4050,  1.5909,  2.6962,
         -3.7314,  1.8437],
        [ 2.5055, -5.4891, -0.0870,  2.2696, -0.5321,  1.7088, -0.7819,  2.1057,
          0.0148, -1.5604],
        [ 3.0821, -3.2848, -1.9754,  3.9027, -0.0772, -4.1832,  0.9734,  1.1843,
         -0.1100, -0.3706],
        [ 0.6146,  0.4998, -2.1531,  0.6684,  0.8023, -0.9562, -0.3057,  0.2265,
          1.6443, -0.2112],
        [-0.6858,  1.8507, -3.0140,  5.5577, -0.6676, -0.8991,  0.1165, -4.4629,
          1.9674, -0.5112],
        [-1.1948,  5.0787, -6.8997,  5.5244, -4.3149,  0.3481,  0.1376,  2.1641,
         -2.1848,  1.5678],
        [-2.0103,  3.0411, -6.4483,  4.4410,  1.0873,  0.1192,  3.2166, -3.4342,
         -1.1943,  0.6516],
        [-0.0470,  1.6777,  0.9133, -4.9525,  6.9630, -8.1715,  5.6771,  0.5204,
          0.7861, -1.9869],
        [ 0.6680,  0.7351,  2.4422, -2.4921,  0.2128, -0.0582, -0.8677, -0.3135,
         -0.4224,  0.6992]])

检查梯度

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.612 秒)

使用 NumPy 和 SciPy 创建扩展#

无参数示例#

带参数示例#

文档

教程

资源