CMU 10-414 Assignments 实验笔记

前言

本文记录了完成《CMU 10-414/714 Deep Learning System》配套 Assignments 的过程和对应笔记。共有 6 个 hw，循序渐进地从头实现了一个深度学习框架，并利用搭建 DL 中厂常见的网络模型，包括 CNN、RNN、Transformer 等。 环境为 Ubuntu 24 @ WSL2。 由于官方自动评分系统目前不再接受非选课学生注册，因此本代码仅保证能够通过已有测试样例。

资源存档

源码来自官方：Assignments

所有代码均上传至 cmu10-414-assignments: cmu10-414-assignments，如官网撤包，可通过 git 回滚获取原始代码。

hw0

第一个 homework 共需完成 7 个函数，第一个很简单，用于熟悉评测系统，直接从第二个函数开始。

parse_mnist

这个函数签名为：parse_mnist(image_filename, label_filename)，用于读取 MNIST 手写数据集。官网 对数据集格式有详细介绍，直接下拉到 FILE FORMATS FOR THE MNIST DATABASE 这部分即可。

整个数据集分为训练集和测试集，包括数字图像和标签。标签文件内前 8Byte 记录了 magic number 和 number of items，之后按照每个样本占 1Byte 的格式组织。图像文件内前 16Byte 记录了非图像数据，之后按照行优先的顺序按照每个像素占 1Byte 的格式以此排布，每个图片共有 28×28 个像素点。

具体实现中，使用 gzip 库按字节读取数据文件，注意整个数据集需要进行标准化，即将每个像素的灰度值除以 255。完整实现为：

import gzip
import numpy as np

def parse_mnist(image_filename, label_filename):
    image_file_handle = gzip.open(image_filename, 'rb')
    label_file_handle = gzip.open(label_filename, 'rb')

    # Read past the headers
    image_file_handle.read(16)
    label_file_handle.read(8)

    # Read the data
    image_data = image_file_handle.read()
    label_data = label_file_handle.read()

    # Close the files
    image_file_handle.close()
    label_file_handle.close()

    # Process image data
    X = np.frombuffer(image_data, dtype=np.uint8).reshape(-1, 28*28).astype(np.float32)
    X = X / 255.0
    
    # Process label data
    y = np.frombuffer(label_data, dtype=np.uint8)

    return X, y

softmax_loss

这个函数签名为：softmax_loss(Z, y)，需要注意的是它计算的是 softmax 损失，或者说是交叉熵损失，而不是进行 softmax 归一化。

照着公式写两行代码即可，不用再赘述：

def softmax_loss(Z, y):
    rows = np.arange(Z.shape[0])
    return -np.mean(Z[rows, y] - np.log(np.sum(np.exp(Z), axis=1)))

softmax_regression_epoch

这个函数签名为：softmax_regression_epoch(X, y, theta, lr = 0.1, batch=100)，要实现的是 softmax 回归一个 epoch 上的训练过程。

首先计算出总的 batch 数，并进行这么多次的循环。在每个循环内，先从 X 和 y 中取出对应样本，然后根据公式计算即可。这里涉及到将 label 转换为独热编码的一个小技巧：E_batch = np.eye(theta.shape[1])[y_batch]，其它则比较简单：

def softmax_loss(Z, y):
    rows = np.arange(Z.shape[0])
    return -np.mean(Z[rows, y] - np.log(np.sum(np.exp(Z), axis=1)))

nn_epoch

这个函数签名为：nn_epoch(X, y, W1, W2, lr = 0.1, batch=100)，要实现一个双层感知机在一个 epoch 上的训练过程。

跟着公式写代码计算即可，需要注意的两个点：

• ReLU 激活函数可以使用 max 函数进行实现：Z1_batch = np.maximum(X_batch @ W1, 0)
• 除以 batch_size 这一步应该提前在计算 G2 的过程，如果放在最后更新 θθ 这一步，存在精度误差，不能通过测试点。

完整代码为：

def softmax_loss(Z, y):
    rows = np.arange(Z.shape[0])
    return -np.mean(Z[rows, y] - np.log(np.sum(np.exp(Z), axis=1)))

softmax_regression_epoch_cpp

这个函数签名为：void softmax_regression_epoch_cpp(const float *X, const unsigned char *y, float *theta, size_t m, size_t n, size_t k, float lr, size_t batch)，这是一个 softmax 回归在 cpp 上的实现版本。

与 Python 自动处理数组索引越界不同，cpp 版本要分开考虑完整的 batch 和最后一轮不完整的 batch。计算 logits 时，需要使用三轮循环模拟矩阵乘法。cpp 版本的实现可以不写出 EyEy​ 矩阵，梯度计算不用使用矩阵计算，直接使用两层循环，判断 class_idx 是否为正确的 label：softmax[sample_idx * k + class_idx] -= (y[start_idx + sample_idx] == class_idx);。

完整的代码为：

def softmax_regression_epoch(X, y, theta, lr=0.1, batch=100):
    total_batches = (X.shape[0] + batch - 1) // batch
    for i in range(total_batches):
        X_batch = X[i*batch:(i+1)*batch]
        y_batch = y[i*batch:(i+1)*batch]
        E_batch = np.eye(theta.shape[1])[y_batch]
        logits = X_batch @ theta
        Z_batch = np.exp(logits)
        Z_batch /= np.sum(Z_batch, axis=1, keepdims=True)
        gradients = X_batch.T @ (Z_batch - E_batch) / batch
        theta -= lr * gradients

hw0 小结

hw0 理应是在 Lecture 2 课前完成的，初学者看到一堆公式应该会很懵逼，但整个 hw 比较简单，照着公式一步步走就能完成（除了双层感知机中奇怪的精度错误），主要还是用来熟悉 NumPy 和基本的 DL 模型。

hw1

第一个 homework 共有六个小问：正向计算、反向梯度、拓扑排序反向模式自动微分、softmax 损失、双层感知机的 SGD 算法。

Implementing forward & backward computation

前两个小问就放在一起讨论了。第一问是通过 NumPy 的 API 实现一些常用的算子，第二问则是通过第一问的算子实现常用算子的梯度实现。

需要注意的是，notebook 中强调了第一问操作的对象是 NDArray，第二问是 Tensor。前者模拟的事这些算子的低层实现，后者则是通过调用这个算子来实现梯度计算，或者说是将梯度计算封装为另一个算子，这样就可以求梯度看作一个普通运算，进而自动求出梯度的梯度。详细解释请看 Lecture 4。

• PowerScaler
  这个算子作用是对张量逐元素求幂。幂指数作为不可学习的参数，在算子实例化时就固定了，因此不用考虑对幂指数的偏导数。这个很简单，应用幂函数的求导公式即可：

def nn_epoch(X, y, W1, W2, lr=0.1, batch=100):
    total_batches = (X.shape[0] + batch - 1) // batch
    for i in range(total_batches):
        X_batch = X[i*batch:(i+1)*batch]
        y_batch = y[i*batch:(i+1)*batch]
        E_batch = np.eye(W2.shape[1])[y_batch]
        Z1_batch = np.maximum(X_batch @ W1, 0)
        G2_batch = np.exp(Z1_batch @ W2)
        G2_batch /= np.sum(G2_batch, axis=1, keepdims=True)
        G2_batch -= E_batch
        G1_batch = (Z1_batch > 0) * (G2_batch @ W2.T)
        gradients_W1 = X_batch.T @ G1_batch
        gradients_W2 = Z1_batch.T @ G2_batch
        W1 -= lr * gradients_W1
        W2 -= lr * gradients_W2

• EWiseDiv
  这个算子的作用是对张量逐元素求商。梯度计算很简单，即 a/ba/b 分别对 aa 和 bb 求偏导：

• DivScalar
  这个算子的作用是将整个张量同除 scalar，和 PowerScalar 一样，scalar 是不要考虑梯度的：

• MatMul
  这个算子的作用是矩阵乘法。这是这门课程到现在第一个具有挑战性的任务。在计算梯度时，根据课程给出的方法，可以得到如下两个表达式：

但但但是，以上只是理论推导。在实际应用中，存在两个问题：1) 矩阵乘法可能是高维矩阵而非二维矩阵相乘，例如 shape 为 (2, 2, 3, 4) 和 (2, 2, 4, 5) 的两个张量相乘；2) 张量乘法过程可能存在广播的情况，这种情况下的梯度怎么处理。

第一个问题，NumPy 基本都为我们处理好了，只要两个张量的倒数两个维度符合二维矩阵乘法且其余维度（也称为批量维度）相等，或者某个批量维度为 1（会进行广播），它们就可以进行张量乘法运算。

天下没有免费的午餐，自动广播带来便利的同时，也带来了第二个问题。求出的 adjoint 或者说偏导，应该和输入参数的维度一致，但根据公式计算得到的梯度的维度和广播后的维度一样，因此要进行 reduce 操作。

以下是我不严谨且非形式化的 reduce 操作推导：假设矩阵 Am×nAm×n​ 经过广播后是 Ap×n×n′Ap×n×n′​，实际上参与计算的就是这个 A′A′。首先直接假设在计算图上用 A′A′ 替代 AA，当 A′@BA′@B（该节点记为 f(x1,…)f(x1​,…)）的某个输入节点 x1x1​ 需要计算梯度时，就会需要计算张量 ∂f/∂x1∂f/∂x1​ 和张量 A′A′ 求得的偏导之间的乘积。接下来我们把 AA 还原，相对应的，f(x1,…)f(x1​,…) 这个节点计算梯度就要将 pp 维度上的偏导数全部加起来，这体现在 Ap×n×n′Ap×n×n′​ 也是将其 pp 维度上的元素全部加起来，得到 Am×n′Am×n′​。

上面这段描述不太清晰，总而言之就是要将广播出来的维度全部 sum 掉。

NumPy 中广播新增的维度只会放在最前面，因此只需要计算出要 sum 掉维度的个数，然后取前 nn 个维度即可，具体见代码：

• Summation
  这个算子的作用是对张量的指定维度求和。设带求和的张量 XX 的维度为 s1×s2×…×sns1​×s2​×…×sn​，那么求和之后的维度就是移除掉 axesaxes 中指示的维度，形式化表达为：

SUM(Xs1×s2×…×sn,axes)=[∑si∈axesX]{sj∣j∉axes}SUM(Xs1​×s2​×…×sn​​,axes)=[si​∈axes∑​X]{sj​∣j∈/axes}​

假设一个输入为的 shape 为 3×2×4×53×2×4×5，在第 0 和 2 的维度上做 summation，输出的 shape 为 2×52×5。反向传播的过程就是先把 out_grad 扩展到 1×2×1×51×2×1×5，然后广播到输入的 shape。

埋个坑，这部分还没有理解，不知道怎么形式化表达求和运算与并对其求导，误打误撞以下代码通过了测试：

• BroadcastTo
  这个算子的作用是将张量广播到指定的 shape。所谓广播，就是将数据在不存在或者大小为 1 的维度上复制多份，使之与目标 shape 相匹配。

关于广播算子正向和梯度运算的分析，可查看 MatMul 算子，其对广播过程有详细讨论。本算子实现代码为：

• Reshape
  这个算子的作用是将张量重整至指定 shape。反向运算则是将张量重整至输入张量的 shape。其代码实现相当简单：

• Negate
  这个算子作用是将整个张量取相反数，反向运算则是再取一次相反数。其代码实现为：

• Transpose
  这个算子的作用是交换指定的两个轴，如果没指定则默认为最后两个轴。注意，这个算子的行为与 np.transpose 不一致，需要调用 API 是 np.swapaxes。反向运算则是再次交换这两个轴：

Topological sort

这一小问要求实现拓扑排序，涉及的知识点都是数据结构的内容，包括图的拓扑排序、后序遍历和 dfs 算法。

在问题说明中明确要求使用树的后序遍历对算法图求解其拓扑序列，简单来说就是如果本节点存在未访问的子节点（inputs），则先访问子节点，否则访问本节点。所谓访问本节点，就是将其标记为已访问，并将其放入拓扑序列。

结合 dfs 算法，求拓扑序列的代码为：

Implementing reverse mode differentiation

终于开始组装我们的自动微分算法了！核心就是理论课中介绍的反向模式 AD 的算法为代码：

其中有几个注意点：

• autograd.py 文件最后一部分提供了一个助手函数 sum_node_list(node_list)，用于在不创造冗余节点的情况下，对一系列 node 求和，对应伪代码中对 vi‾vi​​ 求和的部分；
• 只有存在输入的节点才要计算梯度，初始 input 节点是没法计算梯度的，要进行判断；
• node.op.gradient 返回值类型为 Tuple | Tensor，要分类处理。node.op.gradient_as_tuple 辅助函数可确保返回类型为 tuple。

在写代码之前，最好复习一遍理论；在 debug 的过程中，可以自己画一下计算图，会有奇效。反向模式 AD 具体实现为：

Softmax loss

本问题先要完成对数函数和指数函数的前向和反向计算，然后再完成 softmax 损失，也就是交叉熵损失函数。

根据说明，这里传入的 y 已经转为了独热编码。具体实现根据说明中的公式一点点写即可，没有要特别说明的：

SGD for a two-layer neural network

最后一问，利用前面的组件，实现一个双层感知机及其随机梯度下降算法。注意事项：

• 这里传入的 y 的值是其 label，需要转为独热编码；
• 一定要仔细看题！在计算两个权重的更新值时，应该使用 NumPy 计算，再转为 Tensor。如果直接使用 Tensor 算子计算，每次更新都会在计算图上新增好几个节点，并指数级增长，这会导致后面一些要 600 多 batch 的测试点要跑十几分钟，实际上只要几秒钟就能跑完。如果你遇到了同样的问题，请再读一遍题目要求。
  代码为：

hw 1 小结

明显感觉到，这个 hw 的强度上来了。由于不太熟悉 NumPy 的运算，中间查了不少资料和别人的实现。感谢 @# xx要努力 的文章 1，不少都是参考他的实现。

最后双层感知机的调试，由于使用了 Tensor 算子来实现，跑了十几分钟，最后才发现题干已经要求使用 NumPy 运算。长了个很大的教训，下次一定好好读题。

hw2

Q1: Weight Initialization

Q1 实现的是几种不同的生成参数初始值的方法，结合 init_basic.py 中的辅助函数，照抄 notebook 中给的公式实现，比较简单。注意把 kwargs 传递给辅助函数，里面有 dtype、device 等信息。

Q2: nn_basic

在 Q2，我们将实现几个最基本的 Module 组件。在 Debug 过程中，我遇到了两个很奇怪问题：

• 所有输入和参数都是 float32 类型，但有一个输出是 float64 类型，导致过不了测试点
• 反向传播中，有一个 node 接收到的 out_grad 的 shape 比该节点的输入的 shape 大，但理论上来说二者应该是一致的
  经过漫长的调试追踪，发现第一个问题是因为在实现 DivScalar 即除法时，如果输入是一个实数而非一个矩阵，numpy 进行除法运算的结果默认为 float64，解决方案是显式调用 np.true_divide 进行除法运算，并使用关键字 dtype='float32' 指定返回值类型。

第二个问题是因为 numpy 中许多运算都会进行自动广播，但是该广播操作对我们的 needle 库是不可见的，也无法添加到计算图中，因此导致了反向传播过程的 shape 不匹配。解决方案是修改修改 Q1 中基础算子的实现，在计算前检查 shape 是否匹配。修改后的 ops_mathematic.py 文件内容为：

万事俱备，接下来可以开始完成 Q2 了。

• Linear
  首先要实现一个线性层，其公式为：

Y=XW+BY=XW+B

注意 weight 和 bias 都是 Parameter 类型，如果定义为 Tensor 类型，会导致后面实现优化器过不了测试点。该模块代码为：

• ReLU
  这个模块很简单，调用 ops.relu 即可。

• Sequential
  这个模块的作用是将多个模块封装进一个模块，由其负责将输入在内部按需计算，并给出最终输出。其实现为：

class Sequential(Module): def init(self, *modules): super().init() self.modules = modules

def forward(self, x: Tensor) -> Tensor:
    ### BEGIN YOUR SOLUTION
    y = x
    for module in self.modules:
        y = module(y)
    return y
    
    ### END YOUR SOLUTION

• SoftmaxLoss
  这里实现其是计算 Softmax 损失的模块，在实现过程中可以调用前面实现的数值稳定版本的 LogSumExp，其公式为：

ℓsoftmax(z,y)=log⁡∑i=1kexp⁡zi−zyℓsoftmax​(z,y)=logi=1∑k​expzi​−zy​​

代码骨架中已经提供了一个将标签转换为度和编码的辅助函数，同时记得求的损失应该是在 batch 上的均值，记得做平均。

• LayerNorm1d
  这是第一个比较有挑战性的模块，其中涉及大量的 reshape 和广播操作，必须对每个变量的形状都了如指掌。注意，可以默认输入的 shape 为 (batch_size, feature_size)。计算公式为：

y=w∘xi−Ex+by=w∘((Var[x]+ϵ)1/2)xi​−E[x]​+b​

根据公式照抄即可，但是要注意中间变量的 shape：

• Flatten
  本模块的作用是保留第一个维度为 batchsize，展平剩下维度。使用 ops.resahpe 实现即可：

• BatchNorm1d
  LayerNorm 是在每一个 batch 内部进行标准化操作，而 BatchNorm 是在每一个 feature 内部进行标准化操作。这就导致了每个样本都会对其他样本的推理结果产生影响，因此在推理时应动态计算均值和方差，以供推理时使用。nn.Module 中有一个 training 字段用于标识是否在训练。

与 LayerNorm 类似，在实现过程中运用了大量 reshape 和广播操作，要留意中间变量的形状。

• Dropout
  Dropout 说白了就是以概率 p 随机丢弃一部分输入，并把剩下的输入进行缩放，以确保下一层的输入期望不变。代码骨架提供了 init.randb 用于生成服从二项分布的布尔序列。代码实现为：

• Residual
  残差模块就是将其它模块的输出和输入的和作为新的输出，实现比较简单：

Q3: Optimizer Implementation

在本问题中，我们将实现优化器模块。优化器模块的作用是根据 loss.backward() 计算出的梯度，更新模型的参数。

需要注意的是，本模块默认启用 l2 正则化或者说 weight decay，因此梯度等于 param.grad + weight_decay * param。

• SGD
  首先要实现的优化器是随机梯度下降，注意在更新参数时要先使用 data 方法创建该参数的副本，以避免计算图越来越大。这里还使用了移动平均来计算梯度，初始值默认为 0。代码实现如下：

• Adam
  没什么好说的，照抄公式就行：

Q4: DataLoader Implementation

在本问题中，我们将实现一些数据处理、Dataset 和 DataLoader 类。Dataset 类用于提供标准接口来访问数据集，DataLoader 类是从数据集读取一个 batch 的迭代器。

• RandomFlipHorizontal
  这个方法是按照概率 p 反转一张图片。注意输入数据的格式是 H*W*C，因此只要使用 np.flip 对 W 轴进行翻转即可。

• RandomCrop
  这个方法是对原图进行随机裁剪。其实现裁剪的流程是：先在上下左右填充 padding 个空白像素，然后根据上下偏移量 shift_y 和左右偏移量 shift_y，在填充图中裁切出与原图大小相同的图片。

• MNISTDataset
  这里要实现针对 MNIST 数据集的 Dataset 子类，作为其子类，要实现三个方法：__init__ 方法初始化图片、标签和数据处理函数、__len__ 返回数据集样本数、__getitem__ 方法获取指定下标的数据集。

要注意的是：1) 使用之前实现的 parse_mnist 方法来解析 MNIST 数据集；2) Dataset 父类提供了 apply_transforms 方法对图片进行处理；3) __getitem__ 方法最好支持以列表指定的多下标以批量读取数据集;4) 图片处理函数接受的数据格式是 H*W*C，但 __getitem__ 返回值的格式应当为 batch_size*n。

代码实现为：

• Dataloader
  Dataloader 类是一个迭代器，也挺简单的，见码知义：

Q5: MLPResNet Implementation

到此为止，我们的 needle 库的各基本组件都实现好了，在本问题中，我们将使用他们拼出 MLP ResNet，并在 MNIST 数据集上进行训练。

• Residual Block
  首先是实现一个残差块，按照下图将这一块块积木拼出来就行：

  

• MLP ResNet
  同样是拼积木，注意这里面有 num_blocks 个 Residual Block。
  

• Epoch
  Epoch 方法用来执行一个 epoch 的训练或者推理，并返回平均错误率或者平均损失，这个函数的逻辑是：实例化损失函数 - 从 DataLoader 获取输入 - 模型推理 - 计算损失 - 重置梯度 - 反向传播 - 更新参数 - 计算错误率。

• Train MNIST
  本方法用于在 MNIST 数据集上训练一个 MLP ResNet，本方法的逻辑是：实例化 Dataset- 实例化 DataLoader- 实例化模型 - 实例化优化器 - 迭代 epoch

hw2 小结

到这里，hw2 就已经完结啦。拖拖拖，拖了一个月才做完，本课程的 test 不是很严格，在 Debug hw2 的过程中发现了不少 hw1 中的错误。遇到问题除了自己调试，也建议参考一下别人的实现，能够提升找到问题所在的效率。

hw3

在本次实验中，我们将构建一个简单的底层库，用于实现 NDArray。之前我们是用 NunPy 来实现，这次我们将手动实现该 CPU 和 GPU 版本的底层库，并且不调用现有的高度优化的矩阵乘法或其他操作代码。

Part 1: Python array operations

第一个部分是通过 Python 代码修改 strides、shape、offset 字段来实现一些操作，由于不涉及底层，使用 Python 来实现这些方法效率已经够高了。

在实现前，先浏览一遍 ndarray.py，其提供大量辅助函数以简化实现过程。

• reshape
  reshape 操作就是按照另一种方式来解析内存中的连续一维数据。代码骨架提供了 NDArray.as_strided 方法将一个 NDArray 转换为指定 shape 和 strides，还有 NDArray.compact_strides 方法根据 shape 生成紧密排列情况下的 strides。

使用以上辅助函数后，reshape 的实现就相当简单：

• permute
  permute 操作指的是对 NDArray 的轴进行重排列，例如原始轴排列的顺序是 BHWC，按照 (0,3,1,2) 方式重排列，得到的轴的顺序是 BCHW。重排后索引为 [i, j, k, l]，则重排前索引为 [i, k, l, j]。假设重排前的 strides 是 m, n, p, q，那么使用重排前索引得到元素下标为 im+kn+lp+jq = im+jq+kn+lp，即重排后索引对应的 strides 是 m, q, n, p，即将原始 strides 按照指定序列重排即可得到重排后对应的 strides。

• broadcast_to
  广播操作很好理解，就是将元素在某些维度上复制，例如 (1, 9, 8, 1) -> (9, 9, 8, 2)，那么广播后索引为 (m, n, p, q) 在原始数组上的索引就是 (0, n, p, 0)，即广播的维度上 strides 置为 0 即可实现该效果。

• getitem
  getitem 用于获取制定索引的元素，并以 NDArray 的形式返回。这里需要注意的是索引都是 slice 对象，代码已完成了对索引的预处理，保证所有的索引都是标准 slice，即其 start、stop、step 属性都存在，且在对应 shape 范围内。

结果的 shape 计算比较简单，计算每个维度上的切片包含几个元素即可。strides 用于根据索引计算索引元素在一维数组中的下标，如果该维度上切片步长不为 1，那相当于每次都要跳过几个元素来访问下个元素，定量计算不难发现，新的 strides 就等于该维度上 slice.step 乘上对应的 strides。

接下来计算 offset，由于切片中存在 start 值，因此如果待访问的索引存在某个维度上索引值小于对应切片上的 start 值的，这个元素不应存在新的 NDArray 上。例如，切片在每个维度上的 start 值为 (2, 3, 4, 5)，那么原始索引 (1, 3, 4, 5) 或者 (2, 3, 4, 1) 都在切片后的首个元素之前，应该被 offset 覆盖。因此，offset 值等于每个维度上的 slice.start 乘上对应的 strides。

Part 2: CPU Backend - Compact and setitem

在本部分中，我们将实现 CPU 版本的 compact 和 setitem，前者用于在内存中创建一份紧密排列的数据副本，后者用于在内存中根据给定的数据赋值。

二者有个共同点，就是涉及到可变循环展开。即，由于给定 NDArray 的维度数量是不确定的，无法通过 n 重循环对数据进行遍历。此处我采用的思路是维护一个索引 (0, 0, 0, ..., 0)，每次手动在最后一位执行 +1 操作，当达到对应维度的 shape 值时则进位，直至最高位也向前进位，说明遍历完毕。

这里我定义了两个辅助函数 bool next_indexvector<int32_t>& index, const std::vector<int32_t>& shape) 和 size_t index_to_offset(const std::vector<int32_t>& index, const std::vector<int32_t>& strides, const size_t offset，分别用于遍历索引和将索引转换为下标。二者实现为：

• compact
  compact 函数只要在预分配内存的 out 上将每个位置的值写入即可。鉴于 out 在内存中是连续的，可以使用 out_index++ 来逐个访问，原始数据则通过上述两个辅助函数进行访问：

• setitem
  setitem 按照是否为标量有两个版本，但都挺简单，利用好两个辅助函数逐个访问对应下标即可：

Part 3: CPU Backend - Elementwise and scalar operations

在本 Part 中，我们将完成一些非常简单的算子的 CPU 版本，本任务主要是用于熟悉在 pybind 中注册 cpp 函数的流程。文档中提到，鼓励使用模板、宏等简化实现。

我没有为每个算子都写一个显式函数声明和定义，我首先实现了 void EwiseOp(const AlignedArray& a, const AlignedArray& b, AlignedArray* out, std::function<scalar_t(scalar_t, scalar_t)> op) 和 void ScalarOp(const AlignedArray& a, scalar_t val, AlignedArray* out, std::function<scalar_t(scalar_t, scalar_t)> op)，分别用于逐元素和统一执行函数 op，通过传入不同的函数 op 可以实现不同的操作。

再通过 lambda 表达式对上面这两个函数部分实例化（柯里化），以便其只接受两个参数 a, b 并在 pybind 中注册。

举个栗子，如果想注册一个按元素乘法，那么完整的代码为：

从外向内看，m.def 用于在 pybind 中注册一个方法，该方法名由第一个参数指定，即 ewise_mul，第二个参数用于指定对应的 cpp 函数，这里可以接受函数指针、匿名函数等。注意，在 python 我们调用 ewise_mul，只传入两个 NDArray，因此我们需要对接受三个参数的 EwiseOp 柯里化，即传入 std::multiplies<scalar_t>() 给 EwiseOp，并将其封装为一个匿名函数。

注册方法的这一步每次都要创建一个匿名函数，有点复杂了，这一步也能抽象为一个宏，即：

上述三个宏，分别用于注册按元素、按标量的双目运算符，和单目运算符在 pybind 中的注册。

应用这些宏，注册所有指定的方法：

注意，其中 std::pow 等有多个重载版本，通过 static_cast 关键字可以指定版本。

Part 4: CPU Backend - Reductions

这里要实现两个归约算子 max 和 sum，为了简化实现，这里只对单个维度进行归约。即便在单个维度上，想要实现归约运算也是相当困难的，因此本任务还进行了简化：在调用归约算子前会将待归约维度重排到最后一个维度上，并在调用结束后自动恢复，因此我们只要实现对最后一个维度的归约运算。

经过一系列简化操作，这两个算子实现起来有点过于简单了：对连续的 reduce_size 个元素进行 max/sum 运算作为输出的新元素即可，最后记得在 pybind 中注册这两个方法：

Part 5: CPU Backend - Matrix multiplication

在本模块中，我们将实现矩阵乘法。

• Matmul
  首先要实现的是三重循环版本的矩阵乘法，外层两个循环依次为 out 的行和列，在开始实现之前，记得对 out 数组进行初始化！

• AlignedDot
  本函数的作用是计算两个 TILE*TILE 的矩阵的矩阵乘法计算结果，并将其加到 out 的对应位置。我们是用三重循环来通过代码实现，而在编译时，其将被优化为向量计算。

• MatmulTiled
  这里通过分块来实现矩阵乘法，分块的原理和分块加速的原因在 Lecture 12 都讲过了，此处不再赘述，笔记在：《CMU 10-414 deep learning system》学习笔记 > Lecture 12。

Part 6: GPU Backend - Compact and setitem

从本 Part 开始，我们要写 CUDA 代码了，第一次接触 CUDA 编程的同学可以看一下这个不到 5 小时的教程 CUDA编程基础入门系列（持续更新）_哔哩哔哩_bilibili，快速入门。

本 Part 中，我们将实现 compact 和 setitem 算子。有了之前实现 CPU 版本的经验，先写一个将逻辑索引转换为物理索引的辅助函数：

CompactKernel 根据文档，其作用是将 a 中逻辑下标为 gid 的数据拷贝到 out[gid] 处，注意判断 gid 是否越界，即：

两个 setitem 算子照猫画虎，比较简单，直接贴代码：

Part 7: CUDA Backend - Elementwise and scalar operations

本 Part 将实现一系列比较简单的单目、双目运算符，重点讲一下如何精简代码。

在 CPU 版本中，我们通过 std::function 动态传入 Op 来实现不同的运算，但在 CUDA 的核函数中是不支持 std 的，因此我们改为通过模板来实现。

分别为逐元素运算和标量运算各写一个模板核函数：

CUDA 核函数中调用的其它函数必须也是核函数或者设备函数，因此我们还要为各个算子封装一个类，并重载 () 运算符，以便实例化上述两个模板核函数：

接下来定义主机端接口，以便注册到 pybind11 中：

上述是双目运算符的实现，接下来实现单目运算符。单目运算符也可以像双目一样通过模板实现，但 copilot 直接生成了对应代码，我也懒得改：

最后，将本文件最后 m.def 开头的代码取消注释，将对应接口注册到 pybind11 中即可。

Part 8: CUDA Backend - Reductions

本 Part 将实现两个规约算子 sum 和 max。

和 CPU 版本一样，待归约的元素在内存中是连续排列的。在 CUDA 中，由每个线程负责一个规约任务，其负责的规约范围为 [gid*size, min(gid*size+size, a_size)]，其中 size 是单个线程负责规约的长度，a_size 是输入数据的长度。

核函数中根据具体的规约算子，计算求和或者最大值即可：

Part 9: CUDA Backend - Matrix multiplication

这是最后一个任务，也是最难的一部分。正如文档中所说，想要实现一个矩阵乘法算子还是挺简单的，让每个线程负责一个结果的计算即可。但，如果想使用 cooperative fetching 和 block shared memory register tiling 技术，尤其是按照理论课中提到的伪代码来实现，则要困难得多。

首先贴出理论课中提到的伪代码：

如上图所示，我们要计算的是两个长度为 N 的方阵之间的乘法，结果矩阵 C 会被分块为 (L,L) 的子矩阵，每个 block 负责计算一个子矩阵。

为了计算这个子矩阵，索引为 block_x, block_y 的 block 需要用到的数据为 A'=A[L*block_x:L*block_x+L,:] 和 B'=B[:,L*block_x:L*block_x+L]。A’ 和 B’ 可能比较大，因此在另一维度上按照长度 S 再次分为 N/S 块，分块后的 shape 分别为 (L,S) 和 (S,L)，二者的矩阵乘法结果的 shape 为 (L,L)，将 N/S 块累加即可得到该 block 负责的子矩阵的结果。

后文将使用矩阵的 shape 来指代该矩阵。

在计算单个 (L,S) 和 (S,L) 的乘法时，每个 block 都会将其对应的数据，即图中 A 和 B 的阴影部分，加载进 block 内线程共享的共享内存中。

通过外积计算单个 (L,S) 和 (S,L) 的乘法，该算法简单说就是从 (L,S) 任取一列，从 (S,L) 中任取一行，进行外积运算。将各种组合方式的外积结果累加，即可实现矩阵乘法。

单个外积运算由 block 内的线程共同完成，如图中所示，每个 thread 负责计算的就是 (V,V) 的更小的矩阵。具体来说，从 (L,S) 任取一列的 shape 为 (L,1)，从 (S,L) 任取一行的 shape 为 (1,L)，对二者按照长度为 V 再次进行分块，即分块为 (V,1) 和 (1,V)shape 的两个矩阵，然后由一个线程负责计算二者的外积，得到 shape 为 (V,V) 的结果。

以上就是理论课伪代码中提到的算法，将其改写为 CUDA 代码时需要考虑各种情况，有如下注意点：

• 理论中提到的需要分块的场景，在实践中可能存在不能完美切分，由余数的情况，需要判断是否越界；
• 每个 block 要计算的结果子矩阵是根据该 block 在 grid 中的位置确定的，每个 thread 要计算的外积的部分是根据其在 block 中的位置确定的；
• 理论中的 S 和 L 在代码中均取值为宏定义常量 TILE 4，V 取值为宏定义常量 V 2。

代码中写了比较详细的注释，这部分比较复杂，难以单纯通过文字讲明白，如有问题欢迎留言一起讨论。

hw3 小结

本 hw 主要内容是各算子 CPU 和 GPU 版本的底层实现，由于是第一次接触 CUDA 代码，在实现 GPU 版本的矩阵乘法的时候花了不少时间 Debug，调试到最后甚至要头疼昏睡过去。好在皇天不负苦心人，灵感一瞬间它就来了，谁懂这柳暗花明又一村的感觉。特别感谢 好友 为我讲解矩阵乘法的实现、大半夜不厌其烦地与我一起调试代码。

hw4

本实验中，首先将实现一些算子，然后分别实现 CNN 和 RNN 网络，并在数据集上进行训练。

Part 1: ND Backend

首先将 src/*、autograd.py、ndarray.py 文件中未实现的方法从之前的 hw 中复制过来，然后在 ops_*.py 中实现之前实现过的 op，大部分只要复制粘贴。

提一下我踩过的坑 2：

• autograd.py 中头文件为如下内容，以保证我们这里使用的后端是根据环境变量 NEEDLE_BACKEND 自动切换的，并且不为 NumPy 后端。

• 在 ndarray.py 中 sum 和 max 规约函数是不支持指定多个轴的，需要修改之以便支持多个轴。

• 在 reshape 之前，要调用 compact
• 在创建 Tensor 时，要确保其与其它数据的 device 相同
• 在 autograd.py 中，有一行代码为 __rsub__ = __sub__，其将 Tensor 的 rsub 方法重定向到了 sub 上，然而减法不具备交换律，该行代码是错误的。需要注释该行，并自行定义 rsub 函数。

然后我们来实现新增的三个 op。

• tanh
  tanh 在我们实现的 backend 中已经有对应的接口了，正向传播直接调用即可。tanh 反向传播公式为：

tanh⁡′(x)=1−tanh⁡2(x)tanh′(x)=1−tanh2(x)

反向传播中直接用 1 减去 node 的平方即可。需要注意，这里有一个上面提到的坑，也就是要自定义 rsub 函数。

• stack
  stack 函数是将多个相同 shape 的 Tensor 堆叠起来，并且会产生一个新的维度。正向传播实现的思路是先分配一个目标 shape 的 Tensor，然后通过赋值运算将他们放到目标位置。这里预分配时 Tensor 需要指定 device 与输入的 Tensor device 一致。反向传播调用逆运算 split。

• split
  split 方法是将指定的一个维度全部拆开，需要注意的是拆开之后的维度不需要 keep dim，也就是要进行一次 reshape 操作，而在 reshape 前是需要显式调用 compact 的。反向传播直接调用 stack 方法即可。

Part 2: CIFAR-10 dataset

在本 Part 中，将完成对 CIFAR-10 数据库的解析。首先从之前的 hw 中复制 python/needle/data/data_transforms.py 和 python/needle/data/data_basic.py 两个文件，并修改 data_basic 中 DataLoader::__next__ 方法为：

在之前 hw 中使用 Tensor.make_const 来实现，但其不会根据当前的 backend 自动切换 cached_data 的数据结构。

CIFAR-10 的数据格式参考 CIFAR-10 and CIFAR-100 datasets，简单来说，按照 batch, channel, height, width 的格式排列。__init__ 方法实现参考网站上已经给出的代码读取数据集，然后进行 reshape 和归一化的操作即可，另外两个方法可以直接写出来。

Part 3: Convolutional neural network

在本 Part 中，我们将首先实现一些算子，然后实现一个 CNN 网络并在 CIFAR 数据集上进行训练。

• pad
  pad 操作逻辑为：首先计算出 out 的 shape，创建一个大小为 shape 的全零 Tensor，然后通过切片将原矩阵赋值到对应位置：

• flip
  很难解释为什么，但是 flip 操作通过负 strides 和正 offset 就可以实现。具体来说，将需要 flip 的维度的 stride 值取负，offset 值等于需要 flip 的维度的 strides 乘 shape-1 然后求和。可以结合代码理解上面这段话：

通过操纵 offset 和 strides 实现 flip 在数学角度应该是可以证明的，此处不表。

• dilate/undilate
  dilate 操作之前没有接触过，但下边的公式很形象：

[1234]⟹[1020000030400000][13​24​]⟹​1030​0000​2040​0000​​

参数 dilation 就是 0 的个数。

这个函数的实现思路与 flip 非常接近，先计算 out 的 shape，然后创建空矩阵，然后通过切片选择目标元素：

dilate 和 undilate 互为逆运算，在计算梯度时互相调用即可。

• conv
  首先处理 padding，不难发现，padding 和 conv 之间具有结合性，即如下两行代码是等价的：

因此，第一步就是将 X 进行 pad，作为新的 X。后面通过 im2col 技术和操作 strides 将 X 和 W 向量化，通过矩阵乘法来实现卷积。上述原理见课程笔记：《CMU 10-414 deep learning system》学习笔记 | 周鑫的个人博客。

反向传播推导见博文：2d 卷积梯度推导与实现 | 周鑫的个人博客

实现 Conv 的代码中使用了较多的 permute 重排操作，如果用 transpose 来实现重排太麻烦了，倒不如直接实现个重排的 TensorOp：

最终实现的代码为：

• nn.Conv
  这里将实现一个卷积层。由如下要求：输入输出的格式为 (N,C,H,W)，padding 应满足当 stride=1 时，输出不缩水，支持 bias 项。

首先修改 Kaming uniform 的实现，使之支持对卷积核的初始化。增加一个逻辑，根据参数 shape 是否为 None，在调用 rand 函数时传入不同的形状即可：

hw4 的代码中，对于 NDArray.sum 的实现有问题，当求和的维度指定为空 tuple 时，其不应该进行求和操作，但原始代码无法正确处理这种情况，需要参数 axis 类型为 list 或者 tuple 的分支进行额外的判断，如果为空 list 或者 tuple，输出等于输入：

万事俱备，卷积层的实现调用上边的函数即可。初始化的部分，根据文档描述初始化好权重和偏执项。对于步长为 1 的卷积，卷积结果会缩水 k-1 行 k-1 列，为了确保 shape 不变，卷积时四周要 pad (k-1)/2，又由于传统上 k 为奇数，因此等价于 pad k/2。

前向传播的部分，首先将 X 重排为 NHWC 的格式，然后加上卷积层。如果由偏执项，则将其广播后再加到结果中，最后将结果重排为 NCHW 格式返回即可。完整代码为：

• ResNet 9
  在实现 TensorOp 的子类时，如果需要初始化 Tensor，一定要指定 device。之前在实现 ReLU 生成 mask 时没有指定 device，将导致反向传播失败，这里对其进行修改：

同样，之前在实现 SoftmaxLoss 生成 one hot 时也没有指定 device，这里需要修改：

此外，还发现在 reshape 操作可能没有调用 compact，这里直接修改其实现，在调用 array_api 前进行 compact 操作：

经过一番小修小补，我们的代码已经相当健壮，足以完成这个 ResNet 9🎉。ResNet 9 网络架构如下所示。写代码的过程中有些漏洞咱也没必要妄自菲薄，毕竟这么厉害的两位大佬也难免有笔误的地方。下图中的 ResNet 9 有一层网络架构写错了，已在原图中指出。

首先来实现 ConvBN，传入的四个参数以此为 channels_in，channels_out，kernel_size 和 stride。hw4 的框架代码中提供了 BatchNorm2d，在拷贝 nn_basic.py 文件时不要直接覆盖。剩余的实现很简单，根据示意图搭积木，运行后哪里报 Not Implemented Error 就补哪里，完整代码为：

很遗憾，上述代码在我的设备上并不能通过 ResNet 9 的测试点，误差为 0.09，远超 tolerance 0.01。但其又能通过后续在 CIFAR 10 训练集上训练 2 epoches 的测试点，且误差为 5e-5，远小于 tolerance 0.01。怀疑前一个测试点数据有问题。

Part 4: Recurrent neural network

• RNN Cell
  RNN cell 似乎没有什么坑，照着文档初始化参数，照着公式进行正向传播：

• RNN
  本节任务是完成一个多层 RNN，即堆叠在一起的 RNN，如下图所示。参数中 num_layers 指定了层数，input_size 指的是最下面那层 RNN 的输入的 x 的 size，除底层之外的 cell 的输入都是前一层的输入，即它们的 input_size = hidden_size
  

由上图，可知每一层的输入都是在变化的，因此考虑维护一个 X_input 列表用于存储当前没计算的 cell 的垂直输入。同样，维护一个 h_input 列表存储当前没计算的 cell 的水平输入。具体来说，当计算的 cell 编号为 hijhij​ 时，其用到的输入为 X_input[i] 和 h_input[j]，同时计算结束后 X_input[j] 和 h_input[j] 都要更新为该节点的输出。

对于这个堆叠在一起的 RNN，可以采用从左往右、从下到上，或者从下到上、从左往右的计算方式。我采用的是先垂直再水平的计算顺序。

模型最后要返回两个变量，一个是最后一层的输出 output，即示意图中的 y 的集合，不难发现最后一层的输出就是最后一层的后一层（假设存在）的垂直输入，即我们一直在维护的 X_input。另一个要返回的变量是最后一列隐藏层，同样，这就是我们一直在维护的水平输入 h_input。水到渠成。

需要注意，Tensor 没有实现 getitem 和 setitem 方法，需要切片存取的时候调用之前实现的 split 和 stack 方法即可。

完整代码为：

Part 5: LSTM

本章节将实现 LSTM，LSTM 和上边的 RNN 逻辑相同，照抄公式，这里直接放出代码：

Part 6: Penn Treebank dataset

• Dictionary
  这个类的作用是构建一个从 word 到 id 双向映射的字典，word2idx 通过读取 dict 来实现，idx2word 通过访问 list 来实现：

• Corpus
  这个类的作用类似于 DataLoader，从文件读取原始数据，通过 Dictionary 将其 tokenize，提供 batchify 将其分割为 batch（这个 batch 指的是输入的 x 中同时存在好几个句子），提供 get_batch 方法将单个句子分割为 batch（这是由于 lstm 的水平深度有限，最多同时接受这么多输入）。

具体实现时参考 docstring 描述即可，由示意图，一目了然。完整代码为：

Part 7: Training a word-level language model

这里有个大坑，ndarray 实现的矩阵乘法不支持批量矩乘，如果由三维矩阵乘二维的情况，需要手动 reshape 再乘，再 reshape 回去。

• Embedding
  这个 Module 的作用是将 token 进行一次线性变换，这个操作涉及到批量矩乘：

• LanguageModel
  搭积木，同样设计批量矩乘：

• epoch_general_ptb
  流程和 hw2 中实现的 epoch 很接近，iter_num = n_batch - seq_len 是因为每条句子长度为 n_batch，按照 seq_len 的滑动窗口加载数据集，同时句子的最后一个词不能作为输入（后面没有输出了）。

如果出现没有实现的异常，就从 hw2 中粘过来。

• train/evaluate ptb
  这里有个坑，这两个函数接受的损失函数传进来的是类，但是当我们要调用前面的 epoch 方法时要将其实例化。

def train_ptb(model, data, seq_len=40, n_epochs=1, optimizer=ndl.optim.SGD,
              lr=4.0, weight_decay=0.0, loss_fn=nn.SoftmaxLoss, clip=None,
              device=None, dtype="float32"):
    np.random.seed(4)
    ### BEGIN YOUR SOLUTION
    for epoch in range(n_epochs):
        avg_acc, avg_loss = epoch_general_ptb(data, model, seq_len, loss_fn(),
                                               optimizer(model.parameters(), lr=lr, weight_decay=weight_decay),
                                               clip=clip, device=device, dtype=dtype)
    return avg_acc, avg_loss
    ### END YOUR SOLUTION

def evaluate_ptb(model, data, seq_len=40, loss_fn=nn.SoftmaxLoss,
                 device=None, dtype="float32"):
    np.random.seed(4)
    ### BEGIN YOUR SOLUTION
    avg_acc, avg_loss = epoch_general_ptb(data, model, seq_len, loss_fn(),
                                           device=device, dtype=dtype)
    return avg_acc, avg_loss
    ### END YOUR SOLUTION

hw4 小结

本节最大的难点在于卷积反向传播的推导，当时推导得头秃了。剩余内容基本都是在搭积木和对之前的实现小修小补，也挺烦躁。

总算是完结了，撒花🎉

hw4_extra

Fine，还有一个实验，继续！

Part 1: Implementing the Multi-Head Attention Activation Layer

这部分将完成一个多头自注意层的正向传播部分。在这个类中提供了一系列辅助函数，记得先浏览一遍。

文档中有两点没有提到：

• self.causal 决定了是否要进行掩码
• self.matmul 计算的是 A@B.T 而不是`A@B

之前实现的 dropout 算子有点问题，没有指定 dtype 和 device，需要修改：

由于输入的 KQV 在已经把“头”作为一个独立维度分离出来了，实现多头自注意力就简单很多，直接当作单头一样抄公式即可：

Part 2 Implementing the Self-Attention Layer with trainable parameters

本部分将实现一个多头自注意力层，包括对 KQV 进行 preNorm、分头、调用之前实现的正向传播代码、合并、线性映射。

首先修改 class Matmul 的实现，使之支持当 A 为 batch 时的 batch matmul 计算：

之前实现的 layerNorm1D 只支持 (batch_size, hddien_size) 的格式，在调用 perNorm 之前要手动进行 reshape，或者直接修改 layerNorm 的实现。

之前实现的 Linear 模块有点问题，当不存在 bias 时仍旧会尝试对其访问，需要修改：

分头行动就是先 reshape 再 permute，这一操作在前面的 hw 中已经出现多次，比较熟练。整体实现比较简单，不到十行代码即可：

Part 3 Implementing a prenorm residual Transformer Layer

本节将完成一个残差 Transformer 层，本层没有难度，纯搭积木。搭积木之前照例对我们的积木块打个补丁，上个 Part 中修改的 Linear 层仍有问题，bias 不支持多 batch 维度，修改为一下内容：

接下来就可以愉快地搭积木啦：

Part 4 Implementing the Transformer model

本部分完成的是一个完整的 Transformer 网络。文档中提到，根据每个词在句子中的序号做一个 embed，所以在初始化时要额外初始化一个 embed 层，在数据进入 Transformer 前把这个 embed 加上去。其余部分搭积木：

由于 ops.matmul 中对于 batch matmul 的坑太多了，之前只修改了正向传播部分，反向传播仍未支持 matmul，最后没能实现在数据集上进行训练 Transformer 网络，略有遗憾。

hw4_extra 小结

hw4_extra 难度相比 hw4 低了很多，毕竟没让我们自己手推 Transformer 的反向传播公式，不然又是一场腥风血雨。

这次是真的完结了，撒花🎉