CMU10-414/714 Deep Learning Systems 第二次作业hw1解析

本周讲解dlsyscourse 的第二次作业，也就是hw1的部分。上周的内容主要是个warm up，熟悉一下环境以及一些简单的代码，本周要开始实现一个基础的自动微分框架needle。同样本周作业也放到了 https://github.com/careywyr/dlsyscourse ，colab的题目也翻译成了中文方便阅读,代码里面也移除了关于mugrade的部分，建议先自己做一遍，这样会更有收获。

关于needle框架的介绍在问题1的前面已经给出了描述，具体要注意的就是下面几个关键类：

Value: 计算图中的一个值，可能是通过对其他Value对象的操作计算得出，或者是常量（叶子）Value对象。我们使用了一个通用类（后续版本中可以专门化为其他数据结构，例如张量Tensor），但目前你主要会通过其子类Tensor进行交互（详见下文）。
Op: 计算图中的一个算子。算子需要在compute()方法中定义其前向传播（即如何基于Value对象的数据计算运算结果），并通过gradient()方法定义其反向传播，用于处理传入的输出梯度。编写这些算子的详细信息将在下文中提供。
Tensor: Value类的一个子类，表示计算图中的一个多维数组输出（张量）。在本次作业以及后续的作业中，你将主要使用Tensor这个Value的子类。我们提供了多个便捷函数（例如运算符重载），这些函数使你能够像正常使用Python中的操作符那样来操作张量，但在实现相应操作之前，这些便捷函数不会正常工作。
TensorOp: Op类的一个子类，专门用于返回张量的运算符。你将在本次作业中实现的所有操作都是这种类型。

那么基于上面的基础类，来完成本次的作业。

问题1：实现前向计算（略）

这一题要在 ops_mathematic 这个文件里实现所有类的 compute() 方法，这个应该是这次作业最简单的一题了，没有太多需要介绍的，就是按照公式使用numpy的方法即可。

问题2：反向传播

为了实现计算图的自动微分，那就要计算反向传播，即将函数的导数和传入的反向的梯度相乘。在题目的介绍中给了逐元素加法EWiseAdd和逐元素乘法EWiseMul运算的反向传播的实现。

对于 $f(x,y) = x + y.$ 它对于x或者y的偏导均是1，因此它的梯度就是1乘以反向传播的梯度，还是其本身。
对于 $f(x,y) = x \circ y$ 。它对于x或者y的偏导分别是另一半，即 y 和 x 。那么反向传播的结果就是 $out_grad * y$ 以及 $out_grad * x$ 。

那么这次作业其实就是完成剩下的那些方法的梯度计算。其实绝大部分都是简单的，只要掌握常见基本函数求导的基本公式即可，主要要说明的是下面几个方法：

transpose
reshape
broadcast
summation
matmul
ReLU

Transpose

对于转置，前向传播的时候是将输入按照axes进行翻转，默认是-1，-2轴，反向传播就是将其翻转回来，所以代码如下：

def gradient(self, out_grad, node):
    if self.axes:
        x, y = self.axes[0], self.axes[1]
    else:
        # 如果 axes 是 None，默认交换最后两个轴
        x, y = -1, -2
    return transpose(out_grad, axes=(x, y))

Reshape

Reshape和转置类似，就是将梯度形状还原：

def gradient(self, out_grad, node):
    ### BEGIN YOUR SOLUTION
    # 将梯度 reshape 成原来的形状
    return reshape(out_grad, node.inputs[0].shape)
    ### END YOUR SOLUTION

Broadcast

广播就有点费劲了，广播前向计算的时候实际上会将张量的维度扩展，比如从(1,n)扩展到(m,n)，那么反向传播的时候，需要将梯度压缩回(1,n)。具体来说，就是对第一个维度进行求和。那么就是：

对于那些被广播的维度，在反向传播时需要对这些维度进行求和。
对于那些没有被广播的维度，梯度会直接传递下去，不需要特殊处理。

def gradient(self, out_grad, node):
    lhs_shape = node.inputs[0].shape  # 原始输入的形状
    output_shape = self.shape  # 广播后的目标形状

    # 如果输入形状和输出形状相同，说明没有发生广播，直接返回 out_grad
    if lhs_shape == output_shape:
        return out_grad

    # 计算输入和输出形状的维度差，处理输入比输出少的维度情况
    ndim_diff = len(output_shape) - len(lhs_shape)

    # 需要缩减的维度
    reduce_axes = []

    # 遍历输出张量的每个维度，找到需要缩减的维度
    for i in range(len(output_shape)):
        # 计算对应的输入维度，如果是补上去的（即左边加 1），则视为 1
        input_dim = lhs_shape[i - ndim_diff] if i >= ndim_diff else 1
        output_dim = output_shape[i]

        # 如果输入维度为 1，而输出维度大于 1，则需要在该轴进行求和
        if input_dim == 1 and output_dim != 1:
            reduce_axes.append(i)

    # 如果 reduce_axes 非空，执行 sum 操作以消除广播的影响
    if reduce_axes:
        # out_grad = array_api.sum(out_grad.numpy(), axis=tuple(reduce_axes))
        out_grad = summation(out_grad, axes=tuple(reduce_axes))

    # reshape 回原始输入的形状
    return reshape(out_grad, lhs_shape)

Sum

对于sum操作的反向传播，首先要考虑sum前向计算的时候带来的后果是什么？它在某些维度上进行了累加，那就意味着这些维度消失了，反向传播的时候需要把这些维度分配回来。

def gradient(self, out_grad, node):
    lhs = node.inputs[0]  
    input_shape = lhs.shape  # 原始张量的形状

    if out_grad.shape == ():
        out_grad = array_api.full_like(lhs.numpy(), out_grad.numpy())

    if self.axes is None:
        # 如果没有指定轴，意味着我们对整个张量求和，结果是一个标量
        return broadcast_to(ndl.Tensor(out_grad), input_shape)

    out_grad = out_grad.numpy() if isinstance(out_grad, Tensor) else out_grad

    # 为 out_grad 恢复被求和的维度
    if out_grad.shape != input_shape:
      for axis in sorted(self.axes, reverse=True):
          out_grad = array_api.expand_dims(out_grad, axis)

    # 扩展为原始输入的形状
    result = broadcast_to(ndl.Tensor(out_grad), input_shape)
    return result

MatMul

对于矩阵乘法，其实如果两者的维度是一样的，那么就跟求xy的偏导区别不大，顶多就是注意一下要进行转置，一个是$out_grad y^T $,一个是$ x^T out_grad$。但要考虑一个问题，就是x和y的维度可能不一致，因为不一致并不代表不能进行矩阵乘法，只要x的最后一个维度和y的第一个维度一样即可。所以反向的时候正如上面的broadcast，这里也要sum一下。

def gradient(self, out_grad, node):
    lhs, rhs = node.inputs

    # 始终使用批量矩阵乘法的逻辑
    grad_lhs = matmul(out_grad, transpose(rhs, axes=(-2, -1)))
    grad_rhs = matmul(transpose(lhs, axes=(-2, -1)), out_grad)

    # 处理广播的维度
    if len(lhs.shape) < len(rhs.shape):
        axes_to_sum = tuple(range(len(rhs.shape) - len(lhs.shape)))
        grad_lhs = summation(grad_lhs, axes=axes_to_sum)
    if len(rhs.shape) < len(lhs.shape):
        axes_to_sum = tuple(range(len(lhs.shape) - len(rhs.shape)))
        grad_rhs = summation(grad_rhs, axes=axes_to_sum)

    return grad_lhs, grad_rhs

ReLU

对于ReLU函数，前向传播就是max(0, x)，其实就是个分段函数，那么对于两段的梯度分别是1和0。然后跟out_grad相乘即可。代码实现上可以创建一个新的Tensor方便进行计算：

def gradient(self, out_grad, node):
    lhs = node.inputs[0]
    np_lhs = lhs.numpy()
    # 创建一个掩码：lhs > 0 的位置为 1，其他位置为 0
    grad_mask = array_api.where(np_lhs > 0, 1, 0)
    # 将 out_grad 和 grad_mask 相乘，得到最终的梯度
    return out_grad * ndl.Tensor(grad_mask)

问题3：拓扑排序(Topological Sort)

这一题就是图的排序，不过要确保执行的是后序深度优先搜索（Post-order Depth-First Search, DFS），它的意思是在访问某个节点之前，先递归地访问它的所有子节点或相邻节点，直到所有路径都被遍历完毕，再回溯并记录当前节点。 比如假设有这样一个有向无环图：

    A
   / \
  B   C
 / \   \
D   E   F

通过DFS排序后，结果是DEBFCA。大致逻辑是：

选择一个未访问的起始节点作为起点。
将该节点标记为已访问。
对于该节点的每一个未访问相邻节点，递归执行后序深度优先搜索。
当所有相邻节点都访问完毕后，将当前节点加入遍历结果。

题目中给出了两个待实现的方法，find_topo_sort 和 topo_sort_dfs ，其中 topo_sort_dfs 是用来当做递归辅助的函数。

def find_topo_sort(node_list: List[Value]) -> List[Value]:
    visited = set()  # 用于记录已访问过的节点
    topo_order = []  # 存储拓扑排序的结果

    # 遍历节点列表，执行 DFS
    for node in node_list:
        if node not in visited:
            topo_sort_dfs(node, visited, topo_order)

    # 返回拓扑排序的结果，由于我们在 DFS 中是后序添加节点的，因此需要反转
    return topo_order

def topo_sort_dfs(node, visited, topo_order):
    """Post-order DFS"""
    # 如果当前节点已经访问过，直接返回
    if node in visited:
        return
    visited.add(node)  # 标记当前节点为已访问

    # 遍历所有前驱节点（即指向当前节点的节点），递归执行DFS
    for predecessor in node.inputs:
        if predecessor not in visited:
            topo_sort_dfs(predecessor, visited, topo_order)

    # 在所有前驱节点都被访问后，将当前节点加入到拓扑排序的结果中
    topo_order.append(node)

问题4：实现反向自动微分

这一题需要参考课件里面的关于反向微分计算的伪码：

要实现的方法实际上是要计算神经网络中每个节点（变量）的梯度，即最终输出张量对每个节点的梯度。那么按照图中的描述，大致如下几个步骤:

初始化字典来存储各节点的梯度（Dictionary that records a list of partial adjoints of each node）
反向遍历图中的节点(就是reverse上一题的list)
遍历节点并计算梯度
1. 遍历节点
2. 累加该节点的梯度（Sum up partial adjoints）
3. 计算梯度
4. 遍历inputs将梯度放进去（Propagate gradients to input nodes）

def compute_gradient_of_variables(output_tensor, out_grad):
    """Take gradient of output node with respect to each node in node_list.

    Store the computed result in the grad field of each Variable.
    """
    # Step 1: Dictionary that records a list of partial adjoints of each node
    node_to_grad = {output_tensor: [out_grad]}

    # Step 2: Traverse nodes in reverse topological order (i.e., from output to input)
    reverse_topo_order = list(reversed(find_topo_sort([output_tensor])))

    for node in reverse_topo_order:
        # Step 3: Sum up partial adjoints        
        v_i = sum_node_list(node_to_grad[node])
        # Step 4: Store the computed gradient in the node's grad field
        node.grad = v_i
        if node.is_leaf():
            continue
        gradient = node.op.gradient_as_tuple(v_i, node)
        # Step 5: Propagate gradients to input nodes
        for i, node1 in enumerate(node.inputs):
            if not node_to_grad.get(node1):
              node_to_grad[node1] = []   
            node_to_grad[node1].append(gradient[i])

问题5：计算Softmax Loss

对于要实现方法的输入输出如下：

Z: 一个 2D Tensor，形状为 (batch_size, num_classes)，包含每个类的logit（未归一化的预测值）。
y_one_hot: 一个 2D Tensor，形状为 (batch_size, num_classes)，表示one-hot 编码的真实标签。
返回的是一个标量 Tensor，表示平均的 softmax 损失。

对于损失的数学公式定义为:
$\ell_{\mathrm{softmax}}(z, y) = \log\sum_{i=1}^k \exp z_i - z_y.$
那么我们可以将步骤分解成如下几步：

对每个样本的logits取指数求和
提取每个样本真实标签的logit值
求loss，记得除以batch

def softmax_loss(Z, y_one_hot):
    exp_sum_z = ndl.summation(ndl.exp(Z), axes=(-1,))
    b = Z.shape[0]
    z_y = ndl.summation(ndl.multiply(Z, y_one_hot), axes=(-1,))
    loss = ndl.summation(ndl.log(exp_sum_z) - z_y) / b
    return loss

问题6：两层神经网络的SGD

这一题跟hw0中的nn_epoch类似，不同的是输入的W1和W2是 Tensors，此外上次是手动推导两层网络的反向传播方差，这次实现了自动微分的方法，因此可以直接使用backward() 方法计算梯度。要注意的就是要使用needle里面的方法，要使得参数是Tensor。

def nn_epoch(X, y, W1, W2, lr=0.1, batch=100):
    num_examples, input_dim = X.shape
    _, num_classes = W2.shape
    for i in range(0, num_examples, batch):
        X_batch = X[i:i + batch]
        y_batch = y[i:i + batch]
        x = ndl.Tensor(X_batch)

        y_one_hot = np.zeros((y_batch.size, num_classes))
        y_one_hot[np.arange(y_batch.size), y_batch] = 1
        y_tensor = ndl.Tensor(y_one_hot)

        Z = ndl.matmul(ndl.relu(ndl.matmul(x, W1)), W2)
        loss = softmax_loss(Z, y_tensor)
        loss.backward()

        W1 = ndl.Tensor(W1.numpy() - lr * W1.grad.numpy())
        W2 = ndl.Tensor(W2.numpy() - lr * W2.grad.numpy())
    return W1, W2

本次作业因为很多地方不再是使用NDArray，而是有时候是Tensor，写的时候经常分不清某个变量到底是个啥类型，所以其实我写的代码是有一些模糊的，虽然功能是实现了，但觉得并不是最优雅的写法，因此仅供参考，相信很多同学可以写出更好的方案。