从零实现深度学习框架——实现Debug功能与no_grad

有时候我们编写一个复杂的模型,想知道模型耗时的瓶颈在哪里,或者想知道模型是如何反向传播的。这时候就需要DEBUG功能,本文就来为我们的metagrad实现debug功能。

引言

本着“凡我不能创造的,我就不能理解”的思想,本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架,该框架类似PyTorch能实现自动求导。

要深入理解深度学习,从零开始创建的经验非常重要,从自己可以理解的角度出发,尽量不适用外部完备的框架前提下,实现我们想要的模型。本系列文章的宗旨就是通过这样的过程,让大家切实掌握深度学习底层实现,而不是仅做一个调包侠。
本系列文章首发于微信公众号:JavaNLP

从零实现深度学习框架——实现Debug功能与no_grad
有时候我们编写一个复杂的模型,想知道模型耗时的瓶颈在哪里,或者想知道模型是如何反向传播的。这时候就需要DEBUG功能,本文就来为我们的metagrad实现debug功能。

创建上下文管理器

class Config:
    debug = False

@contextlib.contextmanager
def using_config(name, value):
    # 保存旧值
    old_value = getattr(Config, name)
    # 设置新值
    setattr(Config, name, value)
    try:
        yield
    finally:
        # 最终设回旧值
        setattr(Config, name, old_value)

首先创建一个Config类,它有一个debug属性,用来表示当前是否为DEBUG模式。

contextmanager 这个装饰器(decorator)接收一个生成器(generator),该generator必须只yield一个值出来,该值会被用在with语句中,绑定到as后面的变量。
我们这里只需要修改Config内部状态,不需要返回任何值,可以只加一个yield

创建操作包装类

class OpWrapper:
    '''
    支持反向传播的Debug
    '''

    def __init__(self, name, xs, backward=False):
        self.name = f"back_{name}" if backward else name
        self.xs = xs
        self.output = None

    def __enter__(self):
        if Config.debug:
            self.start = time.time()
        return self

    def __exit__(self, *junk):
        if Config.debug:
            end = (time.time() - self.start) * 1000
            print(
                f"{self.name:>20} : {end:>7.2f} ms {str([y.shape for y in self.xs]):>40} "
                f"{'-> ' + str(self.output.shape) if self.output is not None else ''}"
            )

创建一个操作包装类,实现魔法方法__enter____exit_。当用在with语句中时,会根据Config.debug值来决定是否记录时间,以及打印DEBUG信息。

应用操作包装类

def debug_mode():
    return using_config("debug", True)

首先创建一个函数,修改debugTrue

修改Tensor#backward方法:

    with OpWrapper(t._ctx.__class__.__name__, [t.grad], backward=True):
                # 以逆序计算梯度,调用t相关运算操作的backward静态方法
                # 计算流向其依赖节点上的梯度(流向其下游)
                grads = t._ctx.backward(t._ctx, t.grad.data)

我们只需要将调用backward方法的代码放进OpWrapper的上下文中即可。

测试DEBUG

修改test_sigmoid函数:

def test_sigmoid():
    x = np.array([[0, 1, 2], [0, 2, 4]], np.float32)

    with debug_mode():
        mx = Tensor(x, requires_grad=True)
        y = F.sigmoid(mx)

        tx = torch.tensor(x, requires_grad=True)
        ty = torch.sigmoid(tx)

        assert np.allclose(y.data, ty.data)

        y.sum().backward()
        ty.sum().backward()

        assert np.allclose(mx.grad.data, tx.grad.data)

这里演示了debug_mode的使用,输出如下:

============================= test session starts =============================
collecting ... collected 1 item

test_sigmoid.py::test_sigmoid PASSED                              [100%]            		  
            back_Sum :    0.00 ms                                     [()] 
        back_TrueDiv :    0.00 ms                                 [(2, 3)] 
            back_Add :    0.00 ms                                 [(2, 3)] 
            back_Exp :    0.00 ms                                 [(2, 3)] 
            back_Neg :    0.00 ms                                 [(2, 3)] 


======================== 1 passed, 1 warning in 0.46s =========================

这里以打印出了反向传播中调用的方法、耗时以及操作的维度。

y = F.sigmoid(mx)
y.sum().backward()

首选调用了sigmoid函数,实际上为:

从零实现深度学习框架——实现Debug功能与no_grad

然后为了方便求梯度,我们调用了sum函数。所以反向传播时先经过Sumbackward方法,然后是从零实现深度学习框架——实现Debug功能与no_grad中的除法,再然后是从零实现深度学习框架——实现Debug功能与no_grad中的加法,再是从零实现深度学习框架——实现Debug功能与no_grad,最后是从零实现深度学习框架——实现Debug功能与no_grad

可以看到,整个反向传播过程都打印了出来,而且还有对应的维度,方便我们进行调试。

除此之外,我们还实现类似PyTorch中的no_grad()方法。

实现no_grad

no_grad的意思是,该上下文中的代码不需要计算梯度,常用于推理阶段或者在验证集上验证。

有了上面的工作,我们实现起来就非常简单:

class Config:
    debug = False
    backprop = True  # 是否需要计算并反向传播梯度

首先修改Config增加一个backprop属性,用于判断是否需要计算梯度。

def no_grad():
    return using_config("backprop", False)

然后增加no_grad函数,用于修改backprop属性。表示当前上下文不需要计算梯度。

    def backward(self, grad: "Tensor" = None) -> None:
        '''
        实现Tensor的反向传播
        Args:
            grad: 如果该Tensor不是标量,则需要传递梯度进来

        Returns:

        '''
        # 只能在requires_grad=True的Tensor上调用此方法
        assert self.requires_grad, "called backward on tensor do not require grad"

        if not Config.backprop:
            return

修改backward函数,如果Config.backpropFalse,那么该函数直接返回。

有了此方法,我们以后就可以拆分训练、测试或验证集了。

版权声明:本文为博主愤怒的可乐原创文章,版权归属原作者,如果侵权,请联系我们删除!

原文链接:https://blog.csdn.net/yjw123456/article/details/122590030

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
心中带点小风骚的头像心中带点小风骚普通用户
上一篇 2022年1月19日 下午9:10
下一篇 2022年1月19日

相关推荐