PyTorch 与 NumPy、SciPy 扩展应用入门教程

在深度学习模型构建过程中，时常需借助多种数学工具拓展神经网络功能。本教程讲解如何利用 NumPy 和 SciPy 在 PyTorch 中创建扩展。

一、无参数神经网络层构建

先来看一个无参数的神经网络层示例。我们将创建一个名为 BadFFTFunction 的类，它虽无实际意义，但能帮助理解基础框架。

import torch
from torch.autograd import Function
from numpy.fft import rfft2, irfft2


class BadFFTFunction(Function):
    @staticmethod
    def forward(ctx, input):
        # 将输入张量转换为 NumPy 数组
        numpy_input = input.detach().numpy()
        #输入 对数据进行傅里叶变换并取绝对值
        result = abs(rfft2(numpy_input))
        # 将结果转换回 PyTorch 张量返回
        return input.new(result)

    
    @staticmethod
    def backward(ctx, grad_output):
        # 将梯度输出转换为 NumPy 数组
        numpy_go = grad_output.numpy()
        # 对梯度进行逆傅里叶变换
        result = irfft2(numpy_go)
        # 将结果转换回 PyTorch 张量返回
        return grad_output.new(result)


## 定义函数方便调用该层
def incorrect_fft(input):
    return BadFFTFunction.apply(input)

使用示例：

## 创建一个 8x8 随机张量，设置 requires_grad=True 表示需要计算梯度
input = torch.randn(8, 8, requires_grad=True)
## 通过我们定义的扩展层进行前向传播
result = incorrect_fft(input)
print(result)
## 进行反向传播
result.backward(torch.randn(result.size()))
print(input)

二、具有可学习权重的神经网络层构建

接下来创建一个具有可学习权重的神经网络层，这里涉及 SciPy 的卷积操作。

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):
        # 将输入张量转换为 NumPy 数组
        input_np, filter_np, bias_np = input.detach().numpy(), filter.detach().numpy(), bias.detachnumpy().()
        # 进行互相关操作并加上偏置
        result = correlate2d(input_np, filter_np, mode='valid') + bias_np
        # 保存反向传播所需数据
        ctx.save_for_backward(input, filter, bias)
        # 将结果转换回 PyTorch 张量返回
        return torch.as_tensor(result, dtype=input.dtype)

    
    @staticmethod
    def backward(ctx, grad_output):
        # 获取前向传播保存的数据
        input, filter, bias = ctx.saved_tensors
        # 将梯度输出转换为 NumPy 数组
        grad_output_np = grad_output.detach().numpy()
        # 计算偏置梯度，对所有位置求和
        grad_bias = np.sum(grad_output_np, keepdims=True)
        # 计算输入梯度，进行卷积操作
        grad_input = convolve2d(grad_output_np, filter.numpy(), mode='full')
        # 计算滤波器梯度，进行互相关操作
        grad_filter = correlate2d(input.numpy(), grad_output_np, mode='valid')
        # 将计算的梯度转换回 PyTorch 张量返回
        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)

使用示例：

## 创建一个 3x3 卷积层
module = ScipyConv2d(3, 3)
print("滤波器和偏置：", list(module.parameters()))
## 创建一个 10x10 随机张量，设置 requires_grad=True 表示需要计算梯度
input = torch.randn(10, 10, requires_grad=True)
## 进行前向传播
output = module(input)
print("卷积输出：", output)
## 进行反向传播
output.backward(torch.randn(8, 8))
print("输入梯度：", input.grad)

三、梯度检查

我们可以使用 PyTorch 的 gradcheck 工具来检查我们计算的梯度是否正确。

from torch.autograd.gradcheck import gradcheck


## 创建一个 3x3 卷积层
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("梯度是否正确：", test)

通过本教程，大家可以在编程狮（W3Cschool）平台上轻松掌握 PyTorch 结合 NumPy 和 SciPy 创建扩展层的方法，为构建更复杂的深度学习模型打下坚实基础。后续可在编程狮（W3Cschool）学习更多进阶内容。