从自定义函数到自定义模块：深入理解PyTorch的SquareModule-CFANZ编程社区

大家好，今天我们继续深入学习PyTorch的自定义功能。在上一篇文章中，我们介绍了如何在PyTorch中自定义一个函数，并实现它的前向传播和反向传播。今天，我们将进一步扩展这个概念，看看如何在PyTorch中自定义一个模块。

为了让大家更好地理解，我们会基于之前的平方函数（SquareFunction），创建一个包含偏置的自定义模块（SquareModule）。

回顾：自定义平方函数

在开始之前，让我们快速回顾一下自定义平方函数的实现：

import torch
from torch.autograd import Function

class SquareFunction(Function):
    @staticmethod
    def forward(ctx, *args, **kwargs):
        inp, = args
        ctx.save_for_backward(inp)
        return inp ** 2

    @staticmethod
    def backward(ctx, *grad_output):
        inp, = ctx.saved_tensors
        grad_output, = grad_output
        grad_input = grad_output * 2 * inp
        return grad_input

这个平方函数的前向传播计算输入的平方，反向传播计算输入的梯度。

自定义模块：SquareModule

现在，我们基于这个平方函数，创建一个包含偏置的自定义模块。这个模块将输入平方并加上一个可训练的偏置。

import torch
from torch import nn
from torch.autograd import Function

class SquareFunction(Function):
    @staticmethod
    def forward(ctx, *args, **kwargs):
        inp = args[0] if len(args) > 0 else kwargs['input']
        ctx.save_for_backward(inp)
        return inp ** 2

    @staticmethod
    def backward(ctx, *grad_output):
        inp, = ctx.saved_tensors
        grad_output, = grad_output
        grad_input = grad_output * 2 * inp
        return grad_input

class SquareModule(nn.Module):
    def __init__(self, input_shape, bias=True, *args, **kwargs):
        super().__init__(*args, **kwargs)

        if bias:
            self.bias = nn.Parameter(torch.rand(input_shape))
        else:
            self.register_parameter('bias', None)

    def forward(self, inp):
        return SquareFunction.apply(inp) + self.bias

代码解释

导入必要的模块：我们需要导入torch、nn和Function。
定义SquareFunction类：这个类和之前一样，实现平方的前向和反向传播。
定义SquareModule类：这个类继承自nn.Module，并在初始化时定义一个可训练的偏置。
实现forward方法：在前向传播中，我们调用SquareFunction.apply(inp)计算输入的平方，并加上偏置。

使用自定义模块

定义好自定义模块后，我们可以像使用PyTorch的其他模块一样来使用它。下面是一个简单的例子：

# 使用自定义函数
x = torch.tensor([2.0, 3.0], requires_grad=True)
model = SquareModule(x.shape)
y = model(x)
y.backward(torch.ones_like(x))

print("Input:", x)
print("Output:", y)
print("Input Gradients:", x.grad, model.bias)

代码解释

创建输入张量：我们创建了一个张量x，并设置requires_grad=True，表示我们需要计算它的梯度。
创建模型：我们创建了一个SquareModule实例，并传入输入的形状。
应用自定义模块：我们使用model(x)计算输出。
计算梯度：我们调用y.backward(torch.ones_like(x))来计算梯度。
打印结果：我们打印输入、输出和输入的梯度，以及模型的偏置。

结果分析

运行上面的代码，你会看到如下输出：

Input: tensor([2., 3.], requires_grad=True)
Output: tensor([ 4.2995,  9.9255], grad_fn=<AddBackward0>)
Input Gradients: tensor([ 4.,  6.]) Parameter containing:
tensor([0.2995, 0.9255], requires_grad=True)

输入：x是一个张量，包含两个元素2.0和3.0。
输出：y是x的平方加上偏置，包含两个元素4.2995和9.9255。
输入的梯度：x.grad是输入的梯度，包含两个元素4.0和6.0。这是因为平方函数的导数是2*x，所以2*2.0 = 4.0，2*3.0 = 6.0。
模型的偏置：model.bias是可训练的偏置，包含两个元素0.2995和0.9255。