CMU 10-714 Assignments 实验笔记

CMU 10-714 Assignments 实验笔记

随缘更新ing

前置要求：微积分、线性代数(可能需要一点矩阵论)、DL基础、Python(熟悉 numpy 和 pytorch 更好)。

几乎不使用任何第三方库实现几个深度学习库，包括自动微分、常见神经网络、数据加载器、优化器，并使用 CUDA 实现 GPU 加速。代码量适中，但是基本 pytorch 中有的东西都有了。如果还想深入了解 torch.compiler 相关内容，可以继续学习Machine Learning Compiler，通过 TVM 介绍了 AI Compiler 中的大部分内容。

实验环境为 Arch Linux，环境如下：

既然这门课为 DL Systems, 本文将重点放在实现上，而不是各种数学公式的推导。

由于官方自动评分系统目前不再接受非选课学生注册，因此本代码仅保证能够通过已有测试样例。

本人不是特别熟悉 python 语法，实现过程中可能有不符合 python 规范的地方，欢迎大家指正。

资源存档

课程官网：Deep Learning Systems

实验要求即测试：Assignments

环境配置

采用 conda 进行包管理。 Arch 用户直接 yay 安装 anaconda 即可。

1
2
3
4
source /opt/anaconda/bin/activate base
conda --crate needle python=3.12
conda activate needle
pip install pytest numpy mugrade

HW0

实验仓库地址：Assignment 1

第一个 homework 主要用于简要复习 Machine Learning 的内容。需要实现包括训练两层的线性神经网络在内的 7 个函数。

Basic add function

用来熟悉评测系统，简单的返回 a + b 即可。

1
2
def add(x, y):
    return x + y

然后根据实验提示看看 test 文件内的测试框架，确保无误(谁家好人 a + b 能写错)后运行 python3 -m pytest -k "add" 进行测试。

paese_mnist

读取 MNIST 手写数字数据集。LeCun 老爷子的博客 内对这个数据集的格式作了详细介绍，直接看 FILE FORMATS FOR THE MNIST DATABASE 这个部分即可。

整个数据集分为 4 个文件，分别为 train_image, train_label, test_imgae, test_label。整个数据集直接使用原始的二进制来存储，前 \(2\) 个字节固定为 \(0\)，第 \(3\) 个字节表明数据类型，第 \(4\) 个字节表明维度数。之后的 \(n * 4\) 个字节表明数据的各个维度大小，其中 \(n\) 为将第 \(4\) 个字节视为 uint8 之后的大小。之后即 MNIST 数据集的实际图片。

例如 0x00 0x00 0x08 0x03 60000 28 28 表明，数据集采用 uint8 格式存储(第 \(4\) 个字节为 0x08,约定 0x08 为 uint8)，数据集有 \(3\) 个维度(第 \(4\) 个字节为 0x03)，之后的 \(3 * 4 = 12\) 个字节表明数据集的维度为 (60000 * 28 * 28)。

具体实现中，使用 gzip 库对文件进行 按字节 读取，最后进行标准化即可(每个像素灰度值除以 \(255\))。

1
2
3
4
5
6
7
8
9
10
11
12
13
def parse_mnist(image_filename, label_filename):
    with gzip.open(image_filename, 'rb') as image:
        zero, dtype, dims = struct.unpack('>HBB', image.read(4))
        shape = tuple(struct.unpack('>I', image.read(4))[0] for _ in range(dims))
        image_data = np.frombuffer(image.read(), dtype=np.uint8).reshape(shape[0], -1)
        image_data = image_data.astype(np.float32) / 255.0

    with gzip.open(label_filename, 'rb') as label:
        zero, dtype, dims = struct.unpack('>HBB', label.read(4))
        shape = tuple(struct.unpack('>I', label.read(4))[0] for _ in range(dims))
        label_data = np.frombuffer(label.read(), dtype=np.uint8).reshape(shape[0])

    return (image_data, label_data)

Softmax loss

实现交叉熵损失函数。照着公式写即可，不再赘述： 

1
2
3
4
def softmax_loss(Z, y):
    sum_log = np.log(np.sum(np.exp(Z), axis=1)) # shape (batch_size)
    prob = Z[np.arange(Z.shape[0]), y]
    return np.mean(sum_log - prob)

SGD for Softmax Regression

实现 softmax 回归的一个 epoch 上的训练过程。

首先从数据中每次拿出一个 batch 大小的数据，然后根据公式算即可。label 转换为 one-hot 的小 trick：I_y = np.eye(num_classes)[y_batch]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def softmax_regression_epoch(X, y, theta, lr = 0.1, batch=100):
    m = X.shape[0]   
    num_classes = theta.shape[1]

    for start in range(0, m, batch):
        end = min(start + batch, m)
        X_batch = X[start:end]
        y_batch = y[start:end]
        batch_size = X_batch.shape[0]

        logits = X_batch @ theta
        exp_logits = np.exp(logits)
        Z = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)
        I_y = np.eye(num_classes)[y_batch]

        grad = X_batch.T @ (Z - I_y) / batch_size
        theta -= lr * grad

SGD for Two-layer nn

实现一个双层感知机在一个 epoch 上的训练过程。

ReLu 本质为和 \(0\) 取最大值，张量版本的 max 为 np.maximum。需要注意的是，除法运算能提前就提前，否则可能会导致精度不够。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def nn_epoch(X, y, W1, W2, lr = 0.1, batch=100):
    def ReLU(X):
        return np.maximum(0, X)

    m = X.shape[0]
    num_classes = W2.shape[1]

    for start in range(0, m, batch):
        end = min(start + batch, m)
        X_batch = X[start:end, ]
        y_batch = y[start:end]
        batch_size = X_batch.shape[0]

        Z1 = ReLU(X_batch @ W1)
        Z2_exp = np.exp(Z1 @ W2)
        I_y = np.eye(num_classes)[y_batch]
        G2 = Z2_exp / np.sum(Z2_exp, axis=1, keepdims=True) - I_y
        G1 = G2 @ W2.T * (Z1 > 0).astype(np.float32)

        G_W1 = X_batch.T @ G1 / batch_size
        G_W2 = Z1.T @ G2 / batch_size

        W1 -= lr * G_W1
        W2 -= lr * G_W2

Softmax Regression by cpp

使用 cpp 实现上面的 Softmax 回归。

与 Python 版本不懂，cpp 没有那么多的语法糖，需要手写矩阵乘法，需要处理多维数据映射为一维。不过在处理 ont-hot 方面比 Python 自由了很多。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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)
{
    for (size_t i = 0;i < m;i += batch) {
        assert(batch != 0);
        size_t batch_size = std::min(batch, m - i);
        const float *X_b = X + i * n;
        const unsigned char *y_b = y + i;

        float *logits = new (std::nothrow) float[batch_size * k];
        float *grad_logits = new (std::nothrow) float[batch_size * k];
        float *grad_theta = new (std::nothrow) float[n * k]();

        if (!logits || !grad_logits || !grad_theta) {
            delete[] logits;
            delete[] grad_logits;
            delete[] grad_theta;
            return ;
        }

        // Forward pass: copute logits = X_b * theta
        for (size_t b = 0;b < batch_size;b ++) {
            for (size_t c = 0;c < k;c ++) {
                float tot  = 0.0f;
                for (size_t d = 0;d < n;d ++) {
                    tot += X_b[b * n + d] * theta[d * k + c];
                }
                logits[b * k + c] = tot;
            }
        } 

        // Softmax and loss gradient
        for (size_t b = 0;b < batch_size;b ++) {
            float max_logit = logits[b * k];
            for (size_t c = 1;c < k;c ++) {
                if (logits[b * k + c] > max_logit) {
                    max_logit = logits[b * k + c];
                }
            }

            // grad_logits = exp(x @ theta)
            float sum_exp = 0.0f;
            for (size_t c = 0;c < k;c ++) {
                float exp_val = std::exp(logits[b * k + c] - max_logit);
                grad_logits[b * k + c] = exp_val;
                sum_exp += exp_val;
            }

            // grad_logits = Z
            int true_class = static_cast<int>(y_b[b]);
            for (size_t c = 0;c < k;c ++) {
                grad_logits[b * k + c] /= sum_exp;
            }
            // grad_logits = Z - I_y
            grad_logits[b * k + true_class] -= 1.0f;
        }

        // Backward
        for (size_t b = 0;b < batch_size;b ++) {
            for (size_t d = 0;d < n;d ++) {
                float x_val = X_b[b * n + d];
                for (size_t c = 0;c < k;c ++) {
                    // in-place transpose
                    grad_theta[d * k + c] += x_val * grad_logits[b * k + c];
                }
            }
        }

        float scale = lr / static_cast<float>(batch_size);
        for (size_t j = 0;j < n * k;j ++) {
            theta[j] -= scale * grad_theta[j];
        }

        delete[] logits;
        delete[] grad_logits;
        delete[] grad_theta;
    }  
}

HW0 总结

建议看完 Lecture 2 之后再写这个实验。没经过西瓜书拷打的同学一开始看到满屏幕的公式应该会很懵，不过好在熟悉推导之后实现起来还是比较简单的，很适合用来熟悉 NumPy 和基础 DL 模型。

HW1

这个 homework 共有 6 个部分：算子正向传播，梯度反向传播，拓扑排序、反向自动微分、softmax 损失以及双层感知机的 SGD。

Forward & Backward Computation

需要实现基本算子的前向传播和反向传播。

前两个部分关系比较紧密，这里放在一起说了。实现算子的时候，需要时刻清楚每一步的张量形状。还需要处理 NumPy 中无处不在的 广播。

推导梯度的时候需要注意：当前算子的 out_grad 形状和当前算子的前向运算形状相同，例如令 \(A \in \mathbb{R}^{m \times d}\)，\(B \in \mathbb{R}^{d \times n}\)，当前算子为 Matmul(A, B)，则从后一层传过来的 out_grad 形状为 \(\mathbb{R}^{m \times n}\)。

对应参数的梯度形状和该参数相同。即当 \(X \in \mathbb{R}^{m \times n}\)，则 \(\frac{\partial L}{\partial \mathbf{X}} \in \mathbb{R}^{m \times n}\)。

需要注意的是，前向传播中操作对象是 NDarray，即只进行单纯的数值运算。而反向传播需要扩展计算图，操作对象为 ndl.Tensor。强烈建议通读 needle/autograd.py 中的 class Tensor 的源码，对理解框架很有用。

  • PowerScalar

所有元素对一个标量进行幂运算。输入输出形状相同，张量梯度形式和标量相同。建议所有梯度都返回元组，方便后续写 AD 的时候进行解包。

1
2
3
4
5
6
7
8
9
10
11
class PowerScalar(TensorOp):
    def __init__(self, scalar: int):
        self.scalar = scalar

    def compute(self, a: NDArray) -> NDArray:
        return array_api.power(a, self.scalar)

    def gradient(self, out_grad: Tensor, node: Tensor):
        a = node.inputs[0]
        grad = self.scalar * power_scalar(a, self.scalar - 1)
        return (out_grad * grad, )
  • EWiseDiv

张量对应元素之间的除法。输入输出形状相同，张量梯度形式和标量相同。分别对分子分母求偏导数即可。

1
2
3
4
5
6
7
8
9
class EWiseDiv(TensorOp):
    def compute(self, a, b):
        return array_api.divide(a, b)

    def gradient(self, out_grad, node):
        lhs, rhs = node.inputs
        lhs_grad = power_scalar(rhs, -1)
        rhs_grad = negate(divide(lhs, power_scalar(rhs, 2)))
        return (out_grad * lhs_grad, out_grad * rhs_grad)
  • DivScalar 所有元素对一个标量进行除法。输入输出形状相同，张量梯度形式和标量相同。对分子求导即可。
1
2
3
4
5
6
7
8
9
class EWiseDiv(TensorOp):
    def compute(self, a, b):
        return array_api.divide(a, b)

    def gradient(self, out_grad, node):
        lhs, rhs = node.inputs
        lhs_grad = power_scalar(rhs, -1)
        rhs_grad = negate(divide(lhs, power_scalar(rhs, 2)))
        return (out_grad * lhs_grad, out_grad * rhs_grad)
  • Transpose

广义的矩阵转置，实际更类似 swap 轴的操作。设 \(A \in \mathbb{R}^{2, 3, 4, 5}\), 则 transpose(A, axes=(1, 3, )) 结果为 \(A \in \mathbb{R}^{2, 5, 4, 3}\)，即第 \(1\) 轴和第 \(3\) 个轴交换。当 axes=None 时，默认交换最后两个轴。

因为 \(Y = X^T\)，所以 \(Y_{ji} = X_{ij}\)。对损失 \(L\) 求导：

\[ \frac{\partial L}{\partial X_{ij}}=\sum_{k, l}\frac{\partial L}{\partial Y_{kl}}\cdot\frac{\partial Y_{kl}}{\partial X_{ij}} \]

而 \(\frac{\partial Y_{kl}}{\partial X_{ij}} = 1\) 当且仅当 \(k = j\) 且 \(l = i\)，否则为 \(0\)。

因此有：

\[ \frac{\partial L}{\partial X_{ij}} = \frac{\partial L}{\partial Y_{ji}} \]

即

\[ \frac{\partial L}{\partial X} = {(\frac{\partial L}{\partial Y})}^T \]

所以 transpose 的梯度即为 out_grad 在相同维度上进行一次转置。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Transpose(TensorOp):
    def __init__(self, axes: Optional[tuple] = None):
        self.axes = axes

    def compute(self, a):
        ndim = a.ndim
        if self.axes is None:
            axis1, axis2 = ndim - 2, ndim - 1
        else:
            axis1, axis2 = self.axes
        return array_api.swapaxes(a, axis1, axis2)

    def gradient(self, out_grad, node):
        return (transpose(out_grad, self.axes), )
  • Reshape

本质是重新对内存下标进行映射。

根据链式法则有：

\[ \frac{\partial L}{\partial x_i}=\sum_j\frac{\partial L}{\partial y_j}\cdot\frac{\partial y_j}{\partial x_i} \]

由于 reshape 只是对元素的重排列，某个 \(y_j\) 要么是由某个唯一的 \(x_i\) 直接复制过来，要么和 \(x_i\) 没有任何关系，有 \(\frac{\partial{y_j}}{\partial{x_i}} = 1\)，当且仅当 \(j = i\) 时成立。

即

\[ \frac{\partial L}{\partial x_i}=\frac{\partial L}{\partial y_j} \]

即把 \(\frac{\partial L}{\partial y}\) 按相反的方式进行 reshape 即得到 \(\frac{\partial L}{\partial y}\)。

1
2
3
4
5
6
7
8
9
10
class Reshape(TensorOp):
    def __init__(self, shape):
        self.shape = shape

    def compute(self, a):
        return array_api.reshape(a, self.shape)

    def gradient(self, out_grad, node):
        shape = node.inputs[0].shape
        return (reshape(out_grad, shape), )
  • BoardCast

广播操作本质是沿某个轴进行复制，根据链式法则有：

\[ \frac{\partial L}{\partial x_i}=\sum_{j\in\text{由 }x_i\text{ 复制得到的位置}}\frac{\partial L}{\partial y_j} \cdot \frac{\partial{y_j}}{\partial{x_i}} \]

由于 \(y_j = x_i\)，即 \(\frac{\partial{y_j}}{\partial{x_i}} = 1\)，即沿着广播的轴对 out_grad 进行求和即可。由于 NumPy 中的广播可能会升维，所以最后进行一次 reshape 对齐 \(X\) 的维度即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class BroadcastTo(TensorOp):
    def __init__(self, shape):
        self.shape = shape

    def compute(self, a):
        return array_api.broadcast_to(a, self.shape).copy()

    def gradient(self, out_grad, node):
        input_shape = node.inputs[0].shape
        output_shape = self.shape

        if input_shape == output_shape:
            return (out_grad, )

        ndim_added = len(output_shape) - len(input_shape)
        axes = list(range(ndim_added)) # 添加的维度
        for i, (in_dim, out_dim) in enumerate(zip(input_shape, output_shape[ndim_added:])):
            if in_dim == 1 and out_dim > 1: # 广播的维度
                axes.append(ndim_added + i)

        if axes:
            out_grad = summation(out_grad, tuple(axes))

        out_grad = reshape(out_grad, input_shape)
        return (out_grad,)
  • Summation

这个算子是沿特定轴进行求和，结果进行降维。如果 axes=None，则对全局进行求和，返回标量。

以 \(y = x1 + x2\) 举例，有

\[ \frac{\partial y}{\partial x_1}=1,\quad\frac{\partial y}{\partial x_2}=1 \]

由链式法则有：

\[ \begin{aligned} & \frac{\partial L}{\partial x_1}=\frac{\partial L}{\partial y}\cdot1=\frac{\partial L}{\partial y} \\ & \frac{\partial L}{\partial x_2}=\frac{\partial L}{\partial y}\cdot1=\frac{\partial L}{\partial y} \end{aligned} \]

即 \(y\) 的梯度值对所有参与求和的 \(x\) 值都是相同的。所以只需要将 out_grad 沿着求和的轴广播回去即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Summation(TensorOp):
    def __init__(self, axes: Optional[tuple] = None):
        self.axes = axes

    def compute(self, a):
        return array_api.sum(a, axis=self.axes)

    def gradient(self, out_grad, node):
        input_shape = node.inputs[0].shape

        # 全局求和
        if self.axes is None:
            new_shape = (1, ) * len(input_shape)
            return (broadcast_to(reshape(out_grad, new_shape), input_shape), )

        # 标准化处理，避免 slef.axes 中出现负数轴
        axes = tuple(ax % len(input_shape) if ax < 0 else ax for ax in self.axes)
        new_shape = list(input_shape)
        # 把求和的那几个轴改为1，方便进行广播
        for ax in axes:
            new_shape[ax] = 1

        return (broadcast_to(reshape(out_grad, tuple(new_shape)), input_shape), )
  • Matmul

广义矩阵乘法。NumPy 支持高维矩阵乘法，只需要最后两维度满足矩阵乘法的要求即可，前面的维度都视为 batch。如果维度不匹配，会进行升维或广播。

在计算梯度的时候，根据 Lecture 1中介绍的凑形状方法，可以得到如下两个表达式：

\[ \frac{\partial L}{\partial Y}\frac{\partial Y}{\partial A} = G B^T \\ \frac{\partial L}{\partial Y}\frac{\partial Y}{\partial B} = A^T G \]

但是 NumPy 会进行 reshape 和 boardcast，根据链式法则，需要进行一次规约。沿着广播轴进行求和，并 reshape 为输入的形状即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class MatMul(TensorOp):
    def compute(self, a, b):
        return array_api.matmul(a, b)

    def gradient(self, out_grad, node):
        lhs, rhs = node.inputs
        lhs_grad = matmul(out_grad, transpose(rhs))
        rhs_grad = matmul(transpose(lhs), out_grad)

        def reduce_to_shape(grad, shape):
            if grad.shape == shape:
                return grad
            ndim_added = len(grad.shape) - len(shape)
            axes = list(range(ndim_added))
            for i, (s, gs) in enumerate(zip(shape, grad.shape[ndim_added:])):
                if s == 1 and gs > 1:
                    axes.append(ndim_added + i)
            if axes:
                grad = summation(grad, tuple(axes))
            return reshape(grad, shape)

        return (reduce_to_shape(lhs_grad, lhs.shape),
                reduce_to_shape(rhs_grad, rhs.shape))
  • Negate

按元素取负，梯度即为乘 -1。

1
2
3
4
5
6
7
8
9
10
class Negate(TensorOp):
    def compute(self, a):
        ### BEGIN YOUR SOLUTION
        return array_api.negative(a)
        ### END YOUR SOLUTION

    def gradient(self, out_grad, node):
        ### BEGIN YOUR SOLUTION
        return (negate(out_grad), ) 
        ### END YOUR SOLUTION
  • Log & Exp

这两个都没啥好说的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Log(TensorOp):
    def compute(self, a):
        return array_api.log(a)

    def gradient(self, out_grad, node):
        a = node.inputs[0]
        return (divide(out_grad, a), )

class Exp(TensorOp):
    def compute(self, a):
        return array_api.exp(a)

    def gradient(self, out_grad, node):
        a = node.inputs[0]
        return (out_grad * exp(a), )

Topological Sort

这部分需要实现拓扑排序。文档里叽里呱啦说了半天，说白了就是 后续遍历的 DFS 序即为逆拓扑序。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def find_topo_sort(node_list: List[Value]) -> List[Value]:
    visited = dict()
    topo_order = []
    for node in node_list:
        if not visited.get(node, False):
            topo_sort_dfs(node, visited, topo_order)
    return topo_order

def topo_sort_dfs(node: Value, visited: dict, topo_order: List[Value]) -> None:
    """Post-order DFS"""
    for fa in node.inputs:
        if not visited.get(fa, False):
            topo_sort_dfs(fa, visited, topo_order)
    visited[node] = True
    topo_order.append(node)

有同学会问，为什么不用 BFS 的方式求。答维护每个节点的入度是很大的开销。

reverse mode AD

开始措自动微分了。核心代码和 Lecture 3 中讲的都差不多。autograd.py 中还给我们提供了一个 sum_node_list 函数，对一系列 node 进行求和，对应伪代码中 $=_j 的部分。

由于 ops 中的算子都是继承自 TensorOp，会自动调用 make_from_op，然后拓展计算图。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def compute_gradient_of_variables(output_tensor, out_grad):
    node_to_output_grads_list: Dict[Tensor, List[Tensor]] = {}
    node_to_output_grads_list[output_tensor] = [out_grad]

    reverse_topo_order = list(reversed(find_topo_sort([output_tensor])))

    for node in reverse_topo_order:
        node.grad = sum_node_list(node_to_output_grads_list[node])

        # not a initial/leaf node
        if len(node.inputs) > 0:
            grad = node.op.gradient(node.grad, node)
            for i, fa in enumerate(node.inputs):
                node_to_output_grads_list.setdefault(fa, [])
                node_to_output_grads_list[fa].append(grad[i])

Softmax loss

用 ndl.Tensor 的算子实现 softmax loss。为了方便计算，已经很贴心的把 \(y\) 替换成了独热编码。没啥好说的，照着公式写即可。

1
2
3
4
5
def softmax_loss(Z: ndl.Tensor, y_one_hot: ndl.Tensor):
    sum_log = ndl.log(ndl.summation(ndl.exp(Z), axes=(1, )))
    prob = ndl.summation(ndl.multiply(Z, y_one_hot), axes=(1, ))
    tot_loss = ndl.summation(sum_log - prob)
    return ndl.divide_scalar(tot_loss, Z.shape[0])

SGD for nn

利用前面的逐渐，实现一个双层感知机一个 epoch 的训练过程。需要注意，这里传入的 \(y\) 是 label 类型，需要转换为独热编码。然后在更新权重的时候，建议转换成 NumPy 计算，否则会把不必要的运算加入计算图，导致计算图指数级增长。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def nn_epoch(X, y, W1, W2, lr=0.1, batch=100):
    m = X.shape[0]
    num_classes = W2.shape[1]
    y_one_hot = np.eye(num_classes)[y]

    for start in range(0, m, batch):
        end = min(start + batch, m)
        X_batch = X[start:end, ]
        y_batch = y_one_hot[start:end]
        batch_size = X_batch.shape[0]

        X_tensor = ndl.Tensor(X_batch)
        y_tensor = ndl.Tensor(y_batch)
        Z1 = ndl.relu(ndl.matmul(X_tensor, W1))
        Z2 = ndl.matmul(Z1, W2)
        loss = softmax_loss(Z2, y_tensor)
        loss.backward()

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

    return W1, W2

HW1 总结

强度已经明显大了很多，主要难在推公式上，全靠 GPT 大人推导()。

HW2

Weight Initialization

这四种初始化方法目标一致：让信号在前向传播时每一层的方差都维持在稳定区间，既不会越传越强、也不会逐层衰减，从而避免梯度消失 / 梯度爆炸。实现上就是按 fan_in / fan_out 反推出均匀分布的范围或正态分布的标准差，再交给 init.rand / init.randn 生成张量。

正态分布版本： \[ \text{xavier}: \sigma^2 = \frac{2}{\text{fan\_in}+\text{fan\_out}}, \quad \text{kaiming(relu)}: \sigma^2 = \frac{2}{\text{fan\_in}} \]

均匀分布 \(U(-a, a)\) 的方差是 \(a^2/3\)，反推回去就有 \(a=\sqrt{6/(\text{fan\_in}+\text{fan\_out})}\)（xavier）和 \(a=\sqrt{6/\text{fan\_in}}\)（kaiming）。两者的差别在分母：kaiming 只看 fan_in，是专门为 ReLU 设计的——\(\text{ReLU}\) 会把一半激活压成 0，方差得乘回 2（即 gain \(=\sqrt 2\)）；而反向传播时信号希望 fan_in 和 fan_out 对称，这正是 xavier 分母里多出一个 \(\text{fan\_out}\) 的原因。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def xavier_uniform(fan_in, fan_out, gain=1.0, **kwargs):
    a = gain * math.sqrt(6 / (fan_in + fan_out))
    return rand(fan_in, fan_out, low=-a, high=a, **kwargs)

def xavier_normal(fan_in, fan_out, gain=1.0, **kwargs):
    std = gain * math.sqrt(2 / (fan_in + fan_out))
    return randn(fan_in, fan_out, std=std, **kwargs)

def kaiming_uniform(fan_in, fan_out, nonlinearity="relu", **kwargs):
    assert nonlinearity == "relu", "Only relu supported currently"
    bound = math.sqrt(6 / fan_in)
    return rand(fan_in, fan_out, low=-bound, high=bound, **kwargs)

def kaiming_normal(fan_in, fan_out, nonlinearity="relu", **kwargs):
    assert nonlinearity == "relu", "Only relu supported currently"
    std = math.sqrt(2 / fan_in)
    return randn(fan_in, fan_out, std=std, **kwargs)

nn Basic Module

python/needle/nn/nn_basic.py。这里有条贯穿全程的规矩：Module.parameters() 只认 Parameter 类型（见 _unpack_params），凡是要被优化的张量都得包一层 Parameter(...)，否则优化器根本看不见它。另外 needle 不做隐式广播，w * x + b 之前得手动 broadcast_to。

先看模板里已经写好的骨架（虽然不用自己实现，但决定了所有模块的写法）：parameters() 会递归扫描 self.__dict__ 里的 dict / list / tuple，只要把 Parameter 挂成成员变量就能被挖出来；eval() / train() 也同样递归，会把整棵子树的 training flag 一并改掉。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Parameter(Tensor):
    """A special kind of tensor that represents parameters."""

def _unpack_params(value):
    if isinstance(value, Parameter):
        return [value]
    elif isinstance(value, Module):
        return value.parameters()
    elif isinstance(value, dict):
        params = []
        for k, v in value.items():
            params += _unpack_params(v)
        return params
    elif isinstance(value, (list, tuple)):
        params = []
        for v in value:
            params += _unpack_params(v)
        return params
    else:
        return []

class Module:
    def __init__(self):
        self.training = True

    def parameters(self):
        """Return the list of parameters in the module."""
        return _unpack_params(self.__dict__)

    def _children(self):
        return _child_modules(self.__dict__)

    def eval(self):
        self.training = False
        for m in self._children():
            m.training = False

    def train(self):
        self.training = True
        for m in self._children():
            m.training = True

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)
  • Linear

self.weight 的形状是 (in_features, out_features)，所以 forward 里是 X @ W 而非 PyTorch 的 X @ W.T。注意要先初始化 weight 再初始化 bias——两者都会消耗随机数，顺序一换，后面 mugrade 的期望值就全对不上了。bias 借 kaiming 生成一个 (out_features, 1) 再转置成 (1, out_features)，把标准差控制在合适范围。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Linear(Module):
    def __init__(self, in_features, out_features, bias=True, device=None, dtype="float32"):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features

        weight_tensor = init.kaiming_uniform(in_features, out_features, dtype=dtype, device=device)
        self.weight = Parameter(weight_tensor, requires_grad=True)
        if bias == True:
            bias_tensor = init.kaiming_uniform(out_features, 1, dtype=dtype, device=device)
            self.bias = Parameter(ops.transpose(bias_tensor), requires_grad=True)
        else:
            self.bias = None

    def forward(self, X):
        result = ops.matmul(X, self.weight)
        if self.bias is None:
            return result
        else:
            return result + ops.broadcast_to(self.bias, result.shape)
  • ReLU

HW1 里那个用缓存输入开 mask 的 ops_mathematic.ReLU 直接拿来当模块用即可。

1
2
3
class ReLU(Module):
    def forward(self, x):
        return ops.relu(x)
  • Sequential / Flatten / Residual

Sequential 正着调一遍即可；Flatten 保留 batch 维，把剩下全部展平，本质就是个 reshape；Residual 直接把输入原样加上去。三者都无需多言。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Sequential(Module):
    def __init__(self, *modules):
        super().__init__()
        self.modules = modules

    def forward(self, x):
        out = x
        for module in self.modules:
            out = module(out)
        return out

class Flatten(Module):
    def forward(self, X):
        B = X.shape[0]
        return X.reshape((B, -1))

class Residual(Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x):
        return x + self.fn(x)

LogSumExp 和 LogSoftmax

在 ops/ops_logarithmic.py 里，这是 SoftmaxLoss 的零件，核心在数值稳定性：直接算 \(\exp(1000)\) 会溢出成 inf，\(\exp(-1000)\) 会下溢成 0。做法是先提出最大值 \(M=\max z\)：

\[ \log\sum_i e^{z_i} = \log\left(e^{M}\sum_i e^{z_i-M}\right) = M + \log\sum_i e^{z_i - M} \]

显见所有指数都 \(\le 1\)，从根上避免溢出；那些仍然下溢的项本来就相对最大值可忽略，无伤大雅。参考：https://indii.org/blog/gradients-of-softmax-and-logsumexp/

梯度方面 \(\frac{\partial \text{LSE}(z)}{\partial z_i} = \frac{e^{z_i}}{\sum_j e^{z_j}} = \exp(z_i - \text{LSE}(z))\)，恰好就是 softmax。这里有个坑：compute 里我把归约掉的轴 squeeze 掉了（与 Summation 保持一致），反向传回来的 out_grad 因此少了几维，得先 reshape 成对应轴为 1 的形状再广播回去。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class LogSumExp(TensorOp):
    def __init__(self, axes=None):
        self.axes = axes

    def compute(self, Z):
        max_z = array_api.max(Z, axis=self.axes, keepdims=True)
        sum_z = array_api.sum(array_api.exp(Z - max_z), axis=self.axes, keepdims=True)
        out = array_api.log(sum_z) + max_z
        if self.axes is None:
            return array_api.squeeze(out)
        return array_api.squeeze(out, axis=self.axes)

    def gradient(self, out_grad, node):
        Z = node.inputs[0]
        ndim = len(Z.shape)
        axes = range(ndim) if self.axes is None else self.axes
        shape = tuple([1 if i in axes else Z.shape[i] for i in range(ndim)])

        out_grad_fixed = out_grad.reshape(shape).broadcast_to(Z.shape)
        node_fixed = node.reshape(shape).broadcast_to(Z.shape)
        grad = exp(Z - node_fixed) * out_grad_fixed
        return (grad, )

LogSoftmax 沿 axis=1 归一化，输入按题目约定是 2D 的，\(\text{logsoftmax}(z)_i = z_i - \text{LSE}(z)\)。它的雅可比 \(\partial y_i/\partial z_j = \delta_{ij} - p_j\)（\(p=\text{softmax}(z)\)），所以反向传播恰好是个 rowmatvec：

\[ \text{grad}_z = \text{out\_grad} - p\,\sum_j \text{out\_grad}_j \]

官方签名载的是 LogSoftmax(axes)，但既然已约定只处理 2D + axis=1，我便图省事把 axes 写死在 compute 里。规范做法是放进 __init__。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class LogSoftmax(TensorOp):
    def compute(self, Z):
        axes = (1, )
        max_z = array_api.max(Z, axis=axes, keepdims=True)
        sum_z = array_api.sum(array_api.exp(Z - max_z), axis=axes, keepdims=True)
        return Z - array_api.log(sum_z) - max_z

    def gradient(self, out_grad, node):
        Z = node.inputs[0]
        probs = exp(node)          # node 存的是 logsoftmax，exp 回去就是 softmax

        grad_sum = summation(out_grad, axes=(1, ))
        grad_sum = grad_sum.reshape((Z.shape[0], 1)).broadcast_to(Z.shape)
        grad = out_grad - probs * grad_sum
        return (grad, )

SoftmaxLoss

把 HW1 里那个 loss 用 LogSumExp 重写一遍： \[ \ell_{\text{softmax}}(z,y) = \log\sum_{i=1}^{k}\exp z_i - z_y = -\,\text{logsoftmax}(z)_y \]

也就是拿 logsoftmax 的输出点乘 one-hot，全求和再取负并除以 batch_size。one-hot 现在直接用 init.one_hot 生成即可。

1
2
3
4
5
6
7
class SoftmaxLoss(Module):
    def forward(self, logits, y):
        batch_size, num_classes = logits.shape
        log_probs = ops.logsoftmax(logits)
        y_one_hot = init.one_hot(num_classes, y)
        loss = ops.summation(log_probs * y_one_hot)
        return -(loss / batch_size)

LayerNorm1d

沿 axis=1（特征维）归一化，与 BatchNorm 沿 axis=0（batch 维）恰好相反，因此完全不受 batch 影响，推理时也不需 running stats。方差用题目指定的有偏估计（除 \(N\)，而非 \(N-1\)）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class LayerNorm1d(Module):
    def __init__(self, dim, eps=1e-5, device=None, dtype="float32"):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = Parameter(init.ones(1, dim, device=device, dtype=dtype), device=device, dtype=dtype)
        self.bias = Parameter(init.zeros(1, dim, device=device, dtype=dtype), device=device, dtype=dtype)

    def forward(self, x):
        batch_size, feature_size = x.shape

        mean = ops.summation(x, axes=(1, )) / feature_size
        mean = mean.reshape((batch_size, 1)).broadcast_to(x.shape)

        square = (x - mean) ** 2
        var = ops.summation(square, axes=(1, )) / feature_size
        var = var.reshape((batch_size, 1)).broadcast_to(x.shape)

        weight = self.weight.broadcast_to(x.shape)
        bias = self.bias.broadcast_to(x.shape)
        x_norm = (x - mean) / ops.power_scalar(var + self.eps, 0.5)

        out = weight * x_norm + bias
        return out

weight 全 1、bias 全 0，形状均为 (1, dim)，两者都能直接 broadcast_to(x.shape)；它们同样要包成 Parameter，否则这两条仿射参数不会被优化。均值方差全靠 summation + reshape + broadcast_to 手搓——needle 不做隐式广播，只能这样。

BatchNorm1d

式子与 LayerNorm 完全一致，只是统计量改沿 batch 维计算，并多出一套 running_mean / running_var：训练时用当前 batch 的统计量，同时以 momentum 滑动更新运行估计；推理（self.training == False）时改用运行估计，来一个样本也能正常前向传播（batch_size = 1）。

\[ \hat{x}_{new} = (1-m)\,\hat{x}_{old} + m\,x_{observed} \]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class BatchNorm1d(Module):
    def __init__(self, dim, eps=1e-5, momentum=0.1, device=None, dtype="float32"):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.momentum = momentum

        self.weight = Parameter(init.ones(1, dim, device=device, dtype=dtype), device=device, dtype=dtype)
        self.bias = Parameter(init.zeros(1, dim, device=device, dtype=dtype), device=device, dtype=dtype)

        self.running_mean = Tensor(init.zeros(dim), device=device, dtype=dtype, requires_grad=False)
        self.running_var = Tensor(init.ones(dim), device=device, dtype=dtype, requires_grad=False)

    def forward(self, x):
        # training
        if self.training == True:
            batch_size, feature_size = x.shape
            mean_batch = ops.summation(x, axes=(0, )) / batch_size
            mean = mean_batch.broadcast_to(x.shape)

            square = (x - mean) ** 2
            var_batch = ops.summation(square, axes=(0, )) / batch_size
            var = var_batch.broadcast_to(x.shape)

            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean_batch.data
            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var_batch.data

            weight = self.weight.broadcast_to(x.shape)
            bias = self.bias.broadcast_to(x.shape)
            x_norm = (x - mean) / ops.power_scalar(var + self.eps, 0.5)

            out = weight * x_norm + bias
            return out
        # reasoning
        else:
            weight = self.weight.broadcast_to(x.shape)
            bias = self.bias.broadcast_to(x.shape)
            x_norm = (x - self.running_mean) / ops.power_scalar(self.running_var + self.eps, 0.5)

            out = weight * x_norm + bias
            return out

三个坑：

  1. 更新 running stats 务必用 .data 取裸数组，否则这些值会一直挂在计算图上，图越滚越大，带 memory_check 的测试直接 OOM。
  2. running_mean/running_var 是普通 Tensor(requires_grad=False)，不能包成 Parameter——它们是统计量而不是参数，被 parameters() 抓去交给优化器就完了。方差也必须是这里手搓的有偏估计，直接抄 x.var(axis=0)（PyTorch 默认无偏、除 \(N-1\)）会对不上 mugrade。
  3. train() / eval() 是 Module 递归设置的 flag，藏在 Sequential/Residual 里的 Dropout、BatchNorm 都会跟着切换（test_nn_batchnorm_check_model_eval_switches_training_flag 测的正是这点）。此外，官方文档里 eval 的公式没写仿射项，但 PyTorch 实际是带 \(w, b\) 的，这里照 PyTorch 处理。

Dropout

以概率 \(p\) 置零，再除以 \(1-p\) 把期望拉回来（inverted dropout），这样推理时可直接恒等返回、无需再缩放：

\[ \mathbb{E}[\text{Dropout}(z_j)] = (1-p)\cdot\frac{z_j}{1-p} + p\cdot 0 = z_j \]

1
2
3
4
5
6
7
8
9
10
class Dropout(Module):
    def __init__(self, p=0.5):
        super().__init__()
        self.p = p

    def forward(self, x):
        if not self.training:
            return x
        mask = init.randb(*x.shape, p=1 - self.p)
        return x * mask / (1 - self.p)

init.randb 返回 bool 张量，直接乘上去当 mask 用。注意 \(p\) 是「丢弃」的概率，采样参数应写 \(1-p\)。

优化器

python/needle/optim.py。SGD（带动量和权重衰减）：

\[ u_{t+1} = \beta u_t + (1-\beta)\nabla_\theta f(\theta_t), \qquad \theta_{t+1} = \theta_t - \alpha u_{t+1} \]

1
2
3
4
5
6
7
8
def step(self):
    for param in self.params:
        if param.grad is not None:
            if param not in self.u.keys():
                self.u[param] = ndl.zeros_like(param.data, requires_grad=False)
            self.u[param] = self.momentum * self.u[param] + (1 - self.momentum) * \
                    (param.grad.data + self.weight_decay * param.data)
            param.data = param.data - self.lr * self.u[param]

Adam 再加二阶动量和偏差修正：

\[ \begin{split} u_{t+1} &= \beta_1 u_t + (1-\beta_1) g_t \\ v_{t+1} &= \beta_2 v_t + (1-\beta_2) g_t^2 \\ \hat u_{t+1} &= u_{t+1}/(1-\beta_1^t), \hat v_{t+1} = v_{t+1}/(1-\beta_2^t) \\ \theta_{t+1} &= \theta_t - \alpha\, \hat u_{t+1}/(\hat v_{t+1}^{1/2}+\epsilon) \end{split} \]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def step(self):
    self.t += 1
    for param in self.params:
        if param.grad is not None:
            if param not in self.m.keys():
                self.m[param] = ndl.zeros_like(param.data, requires_grad=False)
            if param not in self.v.keys():
                self.v[param] = ndl.zeros_like(param.data, requires_grad=False)
            grad = param.grad.data + self.weight_decay * param.data
            self.m[param] = self.beta1 * self.m[param].data + (1 - self.beta1) * grad.data
            self.v[param] = self.beta2 * self.v[param].data + (1 - self.beta2) * grad.data * grad.data
            u_hat = self.m[param].data / (1 - self.beta1 ** self.t)
            v_hat = self.v[param].data / (1 - self.beta2 ** self.t)
            param.data = param.data - self.lr * u_hat.data / (ndl.power_scalar(v_hat.data, 0.5) + self.eps).data

要点：

  1. 全程在 .data（裸 NDArray）上算，一点 autograd 都不要碰，否则每次 step() 都在往计算图上接新节点，图一直长，memory_check 测试立刻 OOM。
  2. 不要 in-place 改 param.grad（题目专门强调过）。这里的 weight decay 是加在梯度上的，相当于把 \(\frac{1}{2}\lambda\lVert\theta\rVert^2\) 并进 loss，而非 AdamW 那种 decoupled weight decay。
  3. \(u/m/v\) 这类状态用 dict 按参数存，首次遇到某参数时 zeros_like 开一份；偏差修正用的 \(t\) 是全局步数，在 step 开头自增。忘了 bias correction，前期 \(m, v\) 偏小会把步长带偏。
  4. 题目的动量公式带了 \((1-\beta)\)，与 PyTorch 的 SGD（\(u\leftarrow \beta u + g\)、没有 \(1-\beta\)）不同，照题目来才能对上 mugrade。

数据管道

python/needle/data/ 下面分三块来写。

  • 数据增广在 data_transforms.py：随机水平翻转就是掷骰子决定要不要 np.flip(img, axis=1)（注意轴是 1，即 W）；随机裁剪先 np.pad 补一圈零，再按随机偏移 \(\in[-padding, padding]\) 切回原尺寸。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class RandomFlipHorizontal(Transform):
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img):
        """H x W x C NDArray -> 以概率 p 水平翻转"""
        flip_img = np.random.rand() < self.p
        if flip_img:
            img = np.flip(img, axis=1)
        return img

class RandomCrop(Transform):
    def __init__(self, padding=3):
        self.padding = padding

    def __call__(self, img):
        """零填充后随机裁剪回原尺寸"""
        shift_x, shift_y = np.random.randint(low=-self.padding, high=self.padding+1, size=2)
        img_size = img.shape
        img = np.pad(img, ((self.padding, self.padding), (self.padding, self.padding), (0, 0)), 'constant')
        img = img[self.padding + shift_x:self.padding + shift_x + img_size[0],
                  self.padding + shift_y:self.padding + shift_y + img_size[1], :]
        return img
  • MNISTDataset：__init__ 直接复用 HW1 的 parse_mnist 读入 images/labels；__getitem__ 先 reshape 成 28x28x1 才能喂给 transforms（它们约定输入是 HxWxC），最后返回 (image, label) 二元组。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
class MNISTDataset(Dataset):
    def __init__(self, image_filename, label_filename, transforms=None):
        self.images, self.labels = parse_mnist(image_filename, label_filename)
        self.transforms = transforms

    def __getitem__(self, index):
        x = self.images[index]
        if self.transforms is not None:
            # transforms 要求输入是 H x W x C
            x = self.apply_transforms(x.reshape((28, 28, 1)))
        return x, self.labels[index]

    def __len__(self):
        return self.images.shape[0]

Dataset 基类已经把 apply_transforms（依次跑一遍 transforms 列表）写好了，子类只需要实现 __getitem__ 和 __len__。

  • DataLoader 的 __iter__ / __next__：shuffle=False 时 __init__ 已把 np.arange(len(dataset)) 用 np.array_split 切好；shuffle=True 时须在 __iter__（即每个 epoch 开头）重新 np.random.permutation 一次，并把游标 self.idx 归零——所以 __iter__ 不是摆设，少了它第二个 epoch 就没数据。__next__ 取一段索引、把样本 np.stack 起来，取完 raise StopIteration 让 for 循环正常结束。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class DataLoader:
    def __init__(self, dataset, batch_size=1, shuffle=False):
        self.dataset = dataset
        self.shuffle = shuffle
        self.batch_size = batch_size
        if not self.shuffle:
            self.ordering = np.array_split(np.arange(len(dataset)),
                                           range(batch_size, len(dataset), batch_size))

    def __iter__(self):
        if self.shuffle:
            # reshuffle the data at the beginning of every epoch
            self.ordering = np.array_split(
                np.random.permutation(len(self.dataset)),
                range(self.batch_size, len(self.dataset), self.batch_size),
            )
        self.idx = 0
        return self

    def __next__(self):
        if self.idx >= len(self.ordering):
            raise StopIteration
        batch_indices = self.ordering[self.idx]
        self.idx += 1
        batch = [self.dataset[i] for i in batch_indices]
        # dataset 可能返回单个数组，也可能返回 (image, label) 这种 tuple
        if isinstance(batch[0], tuple):
            return tuple(Tensor(np.stack(col)) for col in zip(*batch))
        return Tensor(np.stack(batch))

文档说 __next__ 返回 NDArray，但这版测试是拿 batch[0].numpy() 取值的，所以要包成 Tensor。zip(*batch) 把「(图, 标签) 的列表」转成「(图堆叠, 标签堆叠)」，顺带解决 label 各自形状 (1,) 对不齐的问题。

MLP ResNet

apps/mlp_resnet.py。残差块为 Linear → norm → ReLU → Dropout → Linear → norm，外面套一层 nn.Residual 再接一个 ReLU：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def ResidualBlock(dim, hidden_dim, norm=nn.BatchNorm1d, drop_prob=0.1):
    return nn.Sequential(
            nn.Residual(
                nn.Sequential(
                    nn.Linear(dim, hidden_dim),
                    norm(hidden_dim),
                    nn.ReLU(),
                    nn.Dropout(drop_prob),
                    nn.Linear(hidden_dim, dim),
                    norm(dim))
                ),
            nn.ReLU()
            )

def MLPResNet(dim, hidden_dim=100, num_blocks=3, num_classes=10,
              norm=nn.BatchNorm1d, drop_prob=0.1):
    return nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.ReLU(),
            *[ResidualBlock(hidden_dim, hidden_dim // 2, norm, drop_prob) for _ in range(num_blocks)],
            nn.Linear(hidden_dim, num_classes),
            )

「Modules should be initialized to match the order of execution」这句并不是建议——每个 nn.Linear 都消耗随机数，构造顺序一变，后面所有层的权重全跟着变，mugrade 期望输出立刻对不上。

epoch 靠 opt 是否给定来切模式：给了就 .train()，没给就 .eval()（测试会故意先调成反方向，就为校验这一点）；loss 和错误率都按样本数加权累加，最后除以总样本数。np.random.seed(4) 是对确定性的要求。train_mnist 就是搭两个 dataloader（训练集 shuffle=True）、建模型、建优化器，再循环 epoch。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def epoch(dataloader, model, opt=None):
    np.random.seed(4)
    if opt is not None:
        model.train()
    else:
        model.eval()
    total_error, total_loss, total_samples = 0.0, 0.0, 0
    loss_fn = nn.SoftmaxLoss()
    for X, y in dataloader:
        batch_size = X.shape[0]
        logits = model(X)
        loss = loss_fn(logits, y)
        if opt is not None:
            opt.reset_grad()
            loss.backward()
            opt.step()
        total_loss += loss.numpy() * batch_size
        pred = logits.numpy().argmax(axis=1)
        total_error += np.sum(pred != y.numpy())
        total_samples += batch_size
    return total_error / total_samples, total_loss / total_samples

默认超参（Adam、lr=1e-3、weight_decay=1e-3、batch=100、hidden=100、3 个 block）跑满 10 epoch，i9-11980HK 上约 7 秒一个 epoch：

epoch train err train loss test err test loss
1 12.11% 0.3975 4.71% 0.1521
2 4.71% 0.1565 4.03% 0.1263
5 2.18% 0.0721 3.08% 0.1011
10 1.49% 0.0485 2.71% 0.0918

测试错误率约 2.7%，比 HW1 手搓的 MLP 好太多。注意后期 train err 仍在降、test err 已基本不降，过拟合的迹象来了。

HW2 总结

比 HW1 更磨人的是各种「隐形约定」而非算法本身：.data 忘了就是拿内存换计算图，Parameter 忘了包就被优化器无视，初始化顺序错一位整套权重就换一套，train/eval flag 决定 BatchNorm 和 Dropout 到底生不生效，需要给每种场景写两套流程。

到这一个基本的 DL System 就已经完成了。接下来就是实现 GPU 后端，以及更复杂的神经网络。

CMU10714
#10714
CMU 10-714 Assignments 实验笔记
https://acmicpc.top/2026/05/23/CMU-10717/
作者
江欣婷
发布于
2026年5月23日
许可协议