torch.autograd.Function 用法及注意事项

智能摘要 AI

本文探讨了如何在 PyTorch 中使用 `torch.autograd.Function` 类自定义操作，并详细分析了一个常见但隐蔽的问题。作者首先介绍了 `Function` 类的用法，展示了如何通过重写 `forward` 和 `backward` 方法来自定义操作，例如实现梯度衰减。然而，作者发现当在前向传播中多次调用不同的衰减系数时，反向传播过程中静态属性未被正确恢复，导致结果不受第一个衰减系数的影响。通过深入分析，作者提出了使用多个 `Function` 子类的解决方案，确保每个衰减系数独立生效。最后，作者强调了编写代码时需仔细验证实际运行过程，避免过度依赖直觉。

众所周知，作为深度学习框架之一的 PyTorch 和其他深度学习框架原理几乎完全一致，都有着自动求导机制，当然也可以说成是自动微分机制。有些时候，我们不想要它自带的求导机制，需要在它的基础之上做些扩展，这个时候我们只需借用 PyTorch 框架中的 Function 类就可以实现了。本文篇幅较长，请诸位读者耐心阅读，最终您会得到您想要的结果。

Function 类的用法

Function 类的用法其实很简单，直接用 PyTorch 官网自带的示例进行说明就绰绰有余。

>>> class Exp(Function):
>>>     @staticmethod
>>>     def forward(ctx, i):
>>>         result = i.exp()
>>>         ctx.save_for_backward(result)
>>>         return result
>>>
>>>     @staticmethod
>>>     def backward(ctx, grad_output):
>>>         result, = ctx.saved_tensors
>>>         return grad_output * result
>>>
>>> # Use it by calling the apply method:
>>> output = Exp.apply(input)

在这里它自定义一个类并继承了 Function 类，然后实现了两个静态的方法。事实上，这两个方法在 Function 类中也有，与其说是实现倒不如说是重写；还有一件事，在 Function 类中的这两个方法被定义为抽象方法，因此与其说是重写倒不如说是必须重写。

顾名思义，foward 方法用来设置前向传播的逻辑，backward 方法用来设置反向传播的逻辑。我们可以发现两个方法有一个参数全部都叫 ctx，这个参数是一个上下文管理器，在调用 forward 的过程中我们可以用该参数的 save_for_backward 方法来保存在 backward 被调用时需要使用的张量。如果需要在 backward 方法中取出保存的张量，可以通过直接访问其属性 saved_tensors 就可以获取到保存的张量。最后一点需要注意的是：foward 方法除去 ctx 参数以外其余参数的个数要与 backward 的返回值数量完全一致，毕竟对于 PyTorch 框架来说每一个输入都有一个梯度，哪怕输入就是一个常数！其余部分非常简单，参数名已经把意思解释得很清楚了，这里就不再赘述了。

观察上面的案例可以发现逻辑其实很简单，它所定义的就是一个把输入传给一个以 e 为底的指数函数的前向传播算法和反向传播算法。

问题重现

在神经网络中，我们有些时候需要在某一层进行反向传播的过程中给梯度乘上一个位于区间 [0, 1) 的常数进行梯度衰减（前向传播按照正常的来，前向传播和反向传播都乘同一个常数 PyTorch 官网有案例）。注意：采用这种梯度衰减作为示例只不过是为了重现之前的我所遇到的问题而已，它目前没有任何理论依据，不要随意使用！这也可以通过类似的方法进行实现，如下所示。

class GradDecay(Function):
    alpha = 0

    @staticmethod
    def forward(ctx, *args, **kwargs):
        return args[0].view_as(args[0])

    @staticmethod
    def backward(ctx, *grad_outputs):
        return GradDecay.alpha*grad_outputs[0]

在这里我通过对类提供一个公有的静态的属性 alpha 来代表乘上的那一个常数。通过 GradDecay.alpha = … 的语句我们可以轻而易举地修改这个属性的值。

在尝试使用它之前我们先定义一个函数，该函数有两个参数，第一个参数表示 forward 方法中的输入，第二个参数代表属性 alpha 的值，代码如下所示。

def grad_decay(x, alpha):
    GradDecay.alpha = alpha
    return GradDecay.apply(x)

接下来我们尝试把它应用到一个简单的多层感知机中，代码如下所示。


class MLP(nn.Module):
    def __init__(self, in_features):
        super(MLP, self).__init__()
        self.linear0 = nn.Linear(in_features, 8)
        self.linear1 = nn.Linear(8, 8)
        self.linear2 = nn.Linear(8, 1)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, alpha=0):
        x = self.relu(self.linear0(x))
        x = grad_decay(x, alpha)
        x = self.relu(self.linear1(x))
        x = grad_decay(x, alpha)
        return self.sigmoid(self.linear2(x))

在这里我定义了一个 4 层的神经网络，我们发现输出是 1，且输出位于区间 (0, 1) 之间，所以很明显，该网络可以用来做目标值在 (0, 1) 之间取任意实数的回归任务，也可以做二分类任务。具体做什么取决于损失函数的定义，一般情况下，如果定义成交叉熵损失函数就是做分类任务，如果定义成平方损失就是做回归任务。

问题马上就来！在这里我调用了两次 grad_decay 函数，我们发现两次调用的时候 alpha 参数是一样的，我们可以把它修改成不一样的吗？如果只看代码会很容易想到一种修改方法，直接去改 forward 方法就行，修改后的 forward 方法如下所示。


def forward(self, x, alpha0=0, alpha1=0):
        x = self.relu(self.linear0(x))
        x = grad_decay(x, alpha0)
        x = self.relu(self.linear1(x))
        x = grad_decay(x, alpha1)
        return self.sigmoid(self.linear2(x))

然而，非常不幸的是只去修改 forward 方法存在非常严重但又极其隐蔽的问题，在说出这个问题是什么的真相之前我首先给出该案例的完整代码，如下所示。

import torch
from sklearn.datasets import load_breast_cancer
from sklearn.metrics import accuracy_score
from torch import nn
from torch.autograd import Function
from torch.optim import Adam


class GradDecay(Function):
    alpha = 0

    @staticmethod
    def forward(ctx, *args, **kwargs):
        return args[0].view_as(args[0])

    @staticmethod
    def backward(ctx, *grad_outputs):
        return GradDecay.alpha*grad_outputs[0]


class MLP(nn.Module):
    def __init__(self, in_features):
        super(MLP, self).__init__()
        self.linear0 = nn.Linear(in_features, 8)
        self.linear1 = nn.Linear(8, 8)
        self.linear2 = nn.Linear(8, 1)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, alpha0=0, alpha1=0):
        x = self.relu(self.linear0(x))
        x = grad_decay(x, alpha0)
        x = self.relu(self.linear1(x))
        x = grad_decay(x, alpha1)
        return self.sigmoid(self.linear2(x))


def grad_decay(x, alpha):
    GradDecay.alpha = alpha
    return GradDecay.apply(x)


def main():
    alpha0, alpha1 = 0.25, 0.5
    bce_loss_func = nn.BCELoss()
    x, y = load_breast_cancer(return_X_y=True)
    x, y = torch.FloatTensor(x), torch.FloatTensor(y)
    torch.manual_seed(0)
    torch.random.manual_seed(0)
    model = MLP(x.shape[1])
    optimizer = Adam(model.parameters())
    for epoch in range(1, 201):
        model.train()
        pred_prob = model(x, alpha0, alpha1).view(-1)
        loss = bce_loss_func(pred_prob, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.eval()
        with torch.no_grad():
            pred_prob = model(x)
        print('epoch: {}, accuracy: {}'.format(epoch, accuracy_score(y.numpy().astype('uint8'), (
                pred_prob > 0.5).view(-1).numpy().astype('uint8'))))


if __name__ == '__main__':
    main()

其运行结果如下图所示：

接下来我先来一针见血的指出哪里有问题，这个非常严重但又极其隐蔽的问题就是：“参数 alpha0 不管修改成多少，结果永远不变！甚至每一个 epoch 的结果和修改之前完全一样！”

出错原因

我们来手工执行核心代码来找出错误原因，首先是多层感知机实例的 forward 方法。

    def forward(self, x, alpha0=0, alpha1=0):
        x = self.relu(self.linear0(x))
        x = grad_decay(x, alpha0)
        x = self.relu(self.linear1(x))
        x = grad_decay(x, alpha1)
        return self.sigmoid(self.linear2(x))


def grad_decay(x, alpha):
    GradDecay.alpha = alpha
    return GradDecay.apply(x)

前向传播方法 forward 的第一行代码没有任何问题，当执行完第二行代码的时候，GradDecay.alpha 的值就变成alpha0（在案例中是 0.25），第三行代码依旧没有任何问题，当执行完第四行代码的时候，GradDecay.alpha 的值就变成了 alpha1（在案例中是 0.5），其他部分也都没有问题。这个 forward 就执行完了，此时需要记住是 GradDecay.alpha 的值就是 0.5！

至于反向传播部分，PyTorch 框架把它全部封装好了，不需要你去单独实现，所以没有代码，但是要想理清反向传播的逻辑只需要把前向传播部分倒过来看。前向传播方法 forward 的最后一行代码进行反向传播没有任何问题，然后把倒数第二行代码进行反向传播的时候需要注意：它不会进入 grad_decay 函数并倒着执行，而是直接转到 GradDecay.backward 方法中去顺着执行（不是倒着执行）！虽然，这一执行过程会让大部分人觉得很懵，但这就是事实！

顺着执行我们发现也没什么毛病，因为 GradDecay.alpha 就等于 alpha1，所以梯度乘上了对应的衰减系数 alpha1，也就是 0.5。

继续回到前向传播方法 forward 中，倒数第三行进行反向传播没有任何问题。倒数第四行需要注意和之前一样，跳到 GradDecay.backward 方法中去顺序执行，此时 GradDecay.alpha 并没有被无缘无故地给修改成 alpha0，也就是 0.25。它依旧是 alpha1，即 0.5。

通过上述的分析过程，这天衣无缝的解释了为什么不管 alpha0 怎么改，结果一点变化都没有。

解决方案

既然找出了问题和出错原因，最后我们就是尝试去解决这个问题，解决这个问题的最容易想到的方法是定义两个 GradDecay 类，一个用来处理第一次调用 grad_decay 函数时所设置的 alpha，即 alpha0，另一个用来处理第一次调用 grad_decay 函数时所设置的 alpha，即 alpha1。修改后的完整代码如下所示。


import torch
from sklearn.datasets import load_breast_cancer
from sklearn.metrics import accuracy_score
from torch import nn
from torch.autograd import Function
from torch.optim import Adam


class GradDecay0(Function):
    alpha = 0

    @staticmethod
    def forward(ctx, *args, **kwargs):
        return args[0].view_as(args[0])

    @staticmethod
    def backward(ctx, *grad_outputs):
        return GradDecay0.alpha*grad_outputs[0]


class GradDecay1(Function):
    alpha = 0

    @staticmethod
    def forward(ctx, *args, **kwargs):
        return args[0].view_as(args[0])

    @staticmethod
    def backward(ctx, *grad_outputs):
        return GradDecay1.alpha*grad_outputs[0]


class MLP(nn.Module):
    def __init__(self, in_features):
        super(MLP, self).__init__()
        self.linear0 = nn.Linear(in_features, 8)
        self.linear1 = nn.Linear(8, 8)
        self.linear2 = nn.Linear(8, 1)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, alpha0=0, alpha1=0):
        x = self.relu(self.linear0(x))
        x = grad_decay(x, alpha0, 0)
        x = self.relu(self.linear1(x))
        x = grad_decay(x, alpha1, 1)
        return self.sigmoid(self.linear2(x))


def grad_decay(x, alpha, a):
    if a:
        GradDecay1.alpha = alpha
        return GradDecay1.apply(x)
    else:
        GradDecay0.alpha = alpha
        return GradDecay0.apply(x)


def main():
    alpha0, alpha1 = 0.25, 0.5
    bce_loss_func = nn.BCELoss()
    x, y = load_breast_cancer(return_X_y=True)
    x, y = torch.FloatTensor(x), torch.FloatTensor(y)
    torch.manual_seed(0)
    torch.random.manual_seed(0)
    model = MLP(x.shape[1])
    optimizer = Adam(model.parameters())
    for epoch in range(1, 201):
        model.train()
        pred_prob = model(x, alpha0, alpha1).view(-1)
        loss = bce_loss_func(pred_prob, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.eval()
        with torch.no_grad():
            pred_prob = model(x)
        print('epoch: {}, accuracy: {}'.format(epoch, accuracy_score(y.numpy().astype('uint8'), (
                pred_prob > 0.5).view(-1).numpy().astype('uint8'))))


if __name__ == '__main__':
    main()

最后需要多说一句的是我之所以没有在 grad_decay 函数中把 else 写成 elif a == 0，是因为我要确保不管在什么条件下都能返回前向传播的输出，而不希望把这个输出变成 None。

在这里 alpha0 是 0.25，这个时候的运行结果如图所示。

为了看出效果，我们尝试把 alpha0 改成一个非常小的数，比如 1e-7，修改后的运行结果如图所示。

虽然还是看不出实质性的效果，但是这可以通过进一步调小 alpha0 来观察到。同时这也至少说明问题已经解决了，不至于出现之前那种不管怎么改 alpha0 每一个 epoch 的每一个 accuracy 都是完全一样。

结论

在这里，我首先通过引入 PyTorch 官方文档所给出的案例介绍了 torch.autograd.Function 类的基本用法，然后通过故意手工实现梯度衰减这荒唐做法来复现我之前遇到的一个非常严重却又极其隐蔽的问题，接着通过手工执行代码的方法找出产生这个问题的原因，最后对症下药给出解决问题的一个方法。最后给出一些注意事项：

进行反向传播的过程中，自定义的 Function 的子类的静态属性不会同时也不可能会被还原！
从想法到代码，一定要仔细的考虑清楚实际运行过程和你所预想的过程是不是完全一致！
不要过度相信自己的感觉，否则当事实与感觉存在偏差的时候会一时半伙无法接受事实！

torch.autograd.Function 用法及注意事项

Function 类的用法

问题重现

出错原因

解决方案

结论

评论 (0)

Eswlnk

随便看看

文章目录

专题展示

WordPress⁵³

工程实践³⁷

torch.autograd.Function 用法及注意事项

Function 类的用法

问题重现

出错原因

解决方案

结论

评论 (0)

猜你喜欢

Eswlnk

随便看看

文章目录

专题展示

WordPress53

工程实践37

热门标签

WordPress⁵³

工程实践³⁷