用 numpy 和 scipy 创建扩展
优质
小牛编辑
137浏览
2023-12-01
译者：cangyunye
作者: Adam Paszke
修订者: Adam Dziedzic
在这个教程里，我们要完成两个任务:
创建一个无参神经网络层。
这里需要调用numpy作为实现的一部分。
创建一个权重自主优化的伸进网络层。
这里需要调用Scipy作为实现的一部分。
import torch
from torch.autograd import Function
无参数示例
这一层并没有特意做什么任何有用的事或者去进行数学上的修正。
它只是被恰当的命名为BadFFTFunction
本层的实现方式
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)
# 由于本层没有任何参数，我们可以简单的声明为一个函数，而不是当做 nn.Module 类
def incorrect_fft(input):
return BadFFTFunction()(input)
Example usage of the created layer:
input = torch.randn(8, 8, requires_grad=True)
result = incorrect_fft(input)
print(result)
result.backward(torch.randn(result.size()))
print(input)
Out:
tensor([[2.2488e-03, 5.1309e+00, 6.4310e+00, 6.0649e+00, 8.1197e+00],
[3.4379e+00, 1.5772e+00, 1.0834e+01, 5.2234e+00, 1.0509e+01],
[2.6480e+00, 1.2934e+01, 9.1619e+00, 1.6011e+01, 9.7914e+00],
[4.0796e+00, 8.6867e+00, 8.8971e+00, 1.0232e+01, 5.7227e+00],
[1.8085e+01, 5.4060e+00, 5.2141e+00, 3.5451e+00, 5.1584e+00],
[4.0796e+00, 8.2662e+00, 1.1570e+01, 8.7164e+00, 5.7227e+00],
[2.6480e+00, 4.5982e+00, 1.1056e+00, 8.8158e+00, 9.7914e+00],
[3.4379e+00, 6.2059e+00, 5.9354e+00, 3.1194e+00, 1.0509e+01]],
grad_fn=<BadFFTFunction>)
tensor([[-0.6461, 0.3270, -1.2190, -0.5480, -1.7273, -0.7326, 0.6294, -0.2311],
[ 0.4305, 1.7503, -0.2914, -0.4237, 0.5441, 1.6597, -0.5645, -0.7901],
[ 0.4248, -2.5986, -0.9257, -0.8651, -0.1673, 1.5749, -1.1857, 1.2867],
[-0.5180, 2.3175, -1.9279, 1.2128, 0.7789, 0.0385, -1.1871, 0.3431],
[ 0.6934, 1.0216, -0.7450, 0.0463, -1.5447, -1.5220, 0.9389, -0.5811],
[ 1.9286, -1.0957, 0.6878, -0.5469, -0.5505, 0.5088, 0.8965, 0.4874],
[-0.2699, 0.3370, 0.3749, -0.3639, -0.0599, 0.8904, 0.1679, -1.8218],
[-0.2963, 0.2246, 0.6617, 1.2258, 0.1530, 0.3114, 0.4568, 0.6181]],
requires_grad=True)
参数化示例
在深度学习的文献中，这一层被意外的称作卷积convolution，尽管实际操作是交叉-关联性cross-correlation的。 (唯一的区别是过滤器filter是为了卷积而翻转，而不是为了交叉关联)。
本层的可自优化权重的实现，依赖于交叉-关联cross-correlation 一个表示权重的滤核 (kernel)。
向后传播计算的和输入相关的梯度以及和过滤器相关的梯度。
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:
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)
Out:
Filter and bias: [Parameter containing:
tensor([[-0.8330, 0.3568, 1.3209],
[-0.5273, -0.9138, -1.0039],
[-1.1179, 1.3722, 1.5137]], requires_grad=True), Parameter containing:
tensor([[0.1973]], requires_grad=True)]
Output from the convolution: tensor([[-0.7304, -3.5437, 2.4701, 1.0625, -1.8347, 3.3246, 2.5547, -1.1341],
[-5.0441, -7.1261, 2.8344, 2.5797, -2.4117, -1.4123, -0.2520, -3.1231],
[ 1.2296, -0.7957, 1.9413, 1.5257, 0.2727, 6.2466, 2.3363, 2.1833],
[-2.6944, -3.3933, 2.3844, 0.2523, -2.0322, -3.1275, -0.2472, 1.5382],
[ 3.6807, -1.1985, -3.9278, 0.8025, 3.3435, 6.6806, 1.1656, 1.3711],
[-1.7426, 1.3875, 8.2674, -0.8234, -4.7534, 3.0932, 1.3048, 2.1184],
[ 0.2095, 1.3225, 0.9022, 3.3324, 0.8768, -5.3459, -1.0970, -4.5304],
[ 2.1688, -1.7967, -0.5568, -9.3585, 0.3259, 5.4264, 2.8449, 6.8120]],
grad_fn=<ScipyConv2dFunctionBackward>)
Gradient for the input map: tensor([[ 7.7001e-01, -2.6786e-02, -1.0917e+00, -4.1148e-01, 2.2833e-01,
-1.7494e+00, -1.4960e+00, 2.3307e-01, 2.2004e+00, 3.1210e+00],
[ 7.0960e-02, 1.8954e+00, 2.0912e+00, -1.3058e+00, -6.1822e-02,
3.8630e+00, -5.1720e-01, -6.9586e+00, -2.5478e+00, -1.4459e+00],
[ 9.3677e-01, -7.5248e-01, 3.0795e-03, -2.1788e+00, -2.6326e+00,
-3.4089e+00, 2.2524e-01, 4.7127e+00, 3.7717e+00, 2.0393e+00],
[-2.0010e+00, 2.7616e+00, 4.0060e+00, -2.0298e+00, 1.6074e+00,
2.3062e+00, -5.4927e+00, -5.3029e+00, 3.5081e+00, 4.5952e+00],
[ 3.4492e-01, -2.3043e+00, -1.5235e+00, -3.3520e+00, -1.3291e-01,
1.4629e+00, 1.9298e+00, 4.5369e-01, -1.5986e+00, -2.3851e+00],
[-2.3929e+00, 5.3965e+00, 5.1353e+00, -1.0269e+00, 2.1031e+00,
-6.2344e+00, -3.6539e+00, -1.7951e+00, -5.6712e-01, 8.6987e-01],
[ 1.1006e-01, -1.5961e+00, 1.2179e+00, 3.4799e-01, -7.1710e-01,
2.5705e+00, 4.5020e-01, 3.8066e+00, 4.8558e+00, 2.1423e+00],
[-9.9457e-01, 1.5614e+00, 1.3985e+00, 3.6700e+00, -1.9708e+00,
-2.4845e+00, 2.5387e+00, -1.2250e+00, -4.6877e+00, -3.3492e+00],
[-4.5289e-01, 2.4210e+00, 3.3681e+00, -2.7785e+00, 1.5472e+00,
-5.0358e-01, -9.7416e-01, 1.1032e+00, 2.0812e-01, 8.2830e-01],
[ 1.1052e+00, -2.5233e+00, 2.0461e+00, 1.1886e-01, -4.8352e+00,
2.4197e-01, -1.5177e-01, -6.9245e-01, -1.8357e+00, -1.5302e+00]])
梯度检查:
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)
Out:
Are the gradients correct: True