这个程序实现了一个完整的深度学习框架，包含以下核心组件：

1. 自动微分系统

• Tensor类：支持梯度追踪和反向传播

• Function类：所有运算的基类，实现前向和反向传播

• 计算图：通过拓扑排序实现自动梯度计算

2. 神经网络层

• Conv2d：卷积层，使用im2col技术优化

• Linear：全连接层

• BatchNorm：批量归一化层

• Dropout：正则化层

• Sequential：顺序容器

3. 损失函数

• CrossEntropyLoss：交叉熵损失（分类任务）

• MSELoss：均方误差损失（回归任务）

4. 优化器

• SGD：支持动量

• Adam：自适应学习率优化器

5. 技术亮点

• 使用im2col/col2im高效实现卷积

• BatchNorm完整的反向传播推导

• 支持广播操作的梯度归约

• 数值稳定的Softmax实现

import numpy as np
from typing import List, Union, Callable, Any, Optional, Tuple

# ============================================================================
# 第一部分：核心张量实现（支持自动微分）
# ============================================================================

class Tensor:
    """
    简易版张量，支持自动微分
    这是整个深度学习框架的基础，类似于PyTorch的Tensor
    
    核心功能：
    1. 存储数据（numpy数组）
    2. 记录计算历史（通过grad_fn）
    3. 支持反向传播（backward方法）
    """
    
    def __init__(self, data: Union[float, list, np.ndarray], requires_grad=False, grad_fn=None, training=True):
        # 将数据统一转为 numpy 数组存储，确保数据类型为float32
        if isinstance(data, (int, float)):
            data = np.array([data], dtype=np.float32)
        elif isinstance(data, list):
            data = np.array(data, dtype=np.float32)
        elif isinstance(data, np.ndarray):
            data = data.astype(np.float32)
        
        self.data = data                # 实际存储的数据
        self.requires_grad = requires_grad  # 是否需要计算梯度
        self.grad = None                # 梯度值，初始为None
        self.grad_fn = grad_fn          # 梯度函数，记录如何计算梯度
        self.training = training        # 训练模式标志（用于Dropout和BatchNorm）
        
    @property
    def shape(self):
        """返回张量形状"""
        return self.data.shape
    
    def __repr__(self):
        return f"Tensor(shape={self.shape}, requires_grad={self.requires_grad})"
    
    def backward(self, gradient=None):
        """
        反向传播，计算所有 requires_grad=True 的 Tensor 的梯度
        使用拓扑排序遍历计算图，从输出到输入依次计算梯度
        
        参数:
            gradient: 初始梯度，如果是标量张量可以省略
        """
        # 如果没有提供初始梯度且当前是标量，则初始梯度设为1
        if gradient is None:
            if self.shape == (1,) or self.shape == ():
                gradient = np.ones_like(self.data)
            else:
                raise RuntimeError("backward需要传入gradient参数，除非是标量张量")
        
        # 拓扑排序：收集所有需要计算梯度的节点
        topo = []           # 存储拓扑排序的结果
        visited = set()     # 记录已访问的节点
        
        def build_topo(tensor):
            """递归构建拓扑排序"""
            if tensor not in visited:
                visited.add(tensor)
                # 如果有梯度函数，先处理它的前驱节点
                if tensor.grad_fn:
                    for prev_tensor in tensor.grad_fn.prev_tensors:
                        if isinstance(prev_tensor, Tensor):
                            build_topo(prev_tensor)
                topo.append(tensor)  # 后序添加
        
        build_topo(self)
        
        # 初始化当前张量的梯度
        self.grad = gradient.copy()
        
        # 反向遍历拓扑序列，从输出到输入
        for tensor in reversed(topo):
            if tensor.grad_fn and tensor.grad is not None:
                # 调用梯度函数计算前驱节点的梯度
                grad_inputs = tensor.grad_fn.backward(tensor.grad)
                if grad_inputs is not None:
                    # 确保返回的是列表或元组
                    if not isinstance(grad_inputs, (list, tuple)):
                        grad_inputs = [grad_inputs]
                    # 将梯度累加到前驱张量上
                    for i, prev_tensor in enumerate(tensor.grad_fn.prev_tensors):
                        if i < len(grad_inputs) and grad_inputs[i] is not None and isinstance(prev_tensor, Tensor):
                            if prev_tensor.requires_grad:
                                if prev_tensor.grad is None:
                                    prev_tensor.grad = grad_inputs[i]
                                else:
                                    prev_tensor.grad += grad_inputs[i]
    
    def zero_grad(self):
        """清零梯度，在每次反向传播前调用"""
        self.grad = None
        # 递归清零所有前驱节点的梯度
        if self.grad_fn:
            for prev in self.grad_fn.prev_tensors:
                if isinstance(prev, Tensor) and prev.requires_grad:
                    prev.zero_grad()
    
    def detach(self):
        """返回一个新的Tensor，不跟踪梯度（类似PyTorch的detach）"""
        return Tensor(self.data.copy(), requires_grad=False)
    
    def numpy(self):
        """转换为numpy数组"""
        return self.data
    
    def item(self):
        """如果是标量，返回Python数值"""
        return self.data.item()
    
    # 以下是一些便捷方法，内部调用对应的Function
    def sum(self, axis=None, keepdims=False):
        return Sum().apply(self, axis, keepdims)
    
    def mean(self, axis=None, keepdims=False):
        return Mean().apply(self, axis, keepdims)
    
    def reshape(self, *shape):
        return Reshape().apply(self, *shape)
    
    def transpose(self, axes=None):
        return Transpose().apply(self, axes)


# ============================================================================
# 第二部分：运算操作与计算图节点
# ============================================================================

class Function:
    """
    运算的基类，所有运算都继承此类
    实现前向传播和反向传播，是自动微分的核心
    
    设计模式：每个运算都是一个Function对象，记录了输入和输出，实现了forward和backward
    """
    
    def __init__(self):
        self.prev_tensors = []   # 前驱张量（输入）
        self.saved_tensors = []  # 保存的中间结果（用于反向传播）
    
    def apply(self, *tensors, **kwargs):
        """
        应用函数到张量，执行前向传播
        参数:
            *tensors: 输入张量
            **kwargs: 额外参数（如axis、keepdims等）
        返回:
            结果张量
        """
        self.prev_tensors = tensors
        # 提取数据并调用forward（如果是Tensor则取data，否则直接使用）
        data_inputs = [t.data if isinstance(t, Tensor) else t for t in tensors]
        result_data = self.forward(*data_inputs, **kwargs)
        
        # 判断结果是否需要梯度：只要任一输入需要梯度，结果就需要梯度
        requires_grad = any(
            (isinstance(t, Tensor) and t.requires_grad) or
            (hasattr(t, 'parameters') and any(isinstance(p, Tensor) and p.requires_grad for p in getattr(t, 'parameters')))
            for t in tensors
        )
        # 创建结果张量，grad_fn指向当前Function
        result = Tensor(result_data, requires_grad=requires_grad, grad_fn=self)
        return result
    
    def save_for_backward(self, *tensors):
        """保存用于反向传播的中间结果"""
        self.saved_tensors = tensors
    
    def forward(self, *inputs):
        """前向传播，子类必须实现"""
        raise NotImplementedError
    
    def backward(self, *grad_output):
        """反向传播，子类必须实现"""
        raise NotImplementedError


def _reduce_grad_to_shape(grad, shape):
    """
    将梯度reduce到指定形状
    用于处理广播操作的反向传播
    """
    gshape = grad.shape
    if gshape == shape:
        return grad
    # 如果梯度维度更多，先求和降维
    if len(gshape) > len(shape):
        for _ in range(len(gshape) - len(shape)):
            grad = np.sum(grad, axis=0)
        gshape = grad.shape
    # 处理维度对齐后的求和
    aligned_shape = (1,) * (len(gshape) - len(shape)) + tuple(shape)
    for i in range(len(gshape)):
        if aligned_shape[i] == 1 and gshape[i] > 1:
            grad = np.sum(grad, axis=i, keepdims=True)
    return grad.reshape(shape)


class Add(Function):
    """加法运算"""
    def forward(self, x, y):
        self.save_for_backward(x.shape, y.shape)
        return x + y
    
    def backward(self, grad_output):
        x_shape, y_shape = self.saved_tensors
        # 梯度需要广播回原始形状
        return _reduce_grad_to_shape(grad_output, x_shape), _reduce_grad_to_shape(grad_output, y_shape)


class Mul(Function):
    """乘法运算"""
    def forward(self, x, y):
        self.save_for_backward(x, y)
        return x * y
    
    def backward(self, grad_output):
        x, y = self.saved_tensors
        grad_x = grad_output * y
        grad_y = grad_output * x
        return _reduce_grad_to_shape(grad_x, x.shape), _reduce_grad_to_shape(grad_y, y.shape)


class MatMul(Function):
    """矩阵乘法运算"""
    def forward(self, x, y):
        self.save_for_backward(x, y)
        return x @ y
    
    def backward(self, grad_output):
        x, y = self.saved_tensors
        # d(x@y)/dx = grad_output @ y.T
        grad_x = grad_output @ y.T
        # d(x@y)/dy = x.T @ grad_output
        grad_y = x.T @ grad_output
        return grad_x, grad_y


class ReLU(Function):
    """ReLU激活函数：max(0, x)"""
    def forward(self, x):
        self.save_for_backward(x)
        return np.maximum(0, x)
    
    def backward(self, grad_output):
        x = self.saved_tensors[0]
        grad = grad_output.copy()
        # 负值处的梯度为0
        grad[x <= 0] = 0
        return grad


class Sum(Function):
    """求和运算"""
    def forward(self, x, axis=None, keepdims=False):
        self.save_for_backward(x, axis, keepdims)
        return np.sum(x, axis=axis, keepdims=keepdims)
    
    def backward(self, grad_output):
        x, axis, keepdims = self.saved_tensors
        # 将梯度的形状还原到输入形状
        if axis is not None:
            if keepdims:
                grad = grad_output
            else:
                # 扩展缺失的维度
                if isinstance(axis, int):
                    grad = np.expand_dims(grad_output, axis=axis)
                else:
                    for ax in sorted(axis):
                        grad_output = np.expand_dims(grad_output, axis=ax)
                    grad = grad_output
            # 广播到原始形状
            grad = np.broadcast_to(grad, x.shape)
        else:
            grad = grad_output * np.ones_like(x)
        return grad


class Mean(Function):
    """均值运算"""
    def forward(self, x, axis=None, keepdims=False):
        self.save_for_backward(x, axis, keepdims)
        return np.mean(x, axis=axis, keepdims=keepdims)
    
    def backward(self, grad_output):
        x, axis, keepdims = self.saved_tensors
        # 计算缩放因子：1/元素个数
        if axis is None:
            scale = 1.0 / x.size
        else:
            if isinstance(axis, int):
                scale = 1.0 / x.shape[axis]
            else:
                size = 1
                for ax in axis:
                    size *= x.shape[ax]
                scale = 1.0 / size
        
        # 还原形状（同Sum）
        if axis is not None:
            if keepdims:
                grad = grad_output
            else:
                if isinstance(axis, int):
                    grad = np.expand_dims(grad_output, axis=axis)
                else:
                    for ax in sorted(axis):
                        grad_output = np.expand_dims(grad_output, axis=ax)
                    grad = grad_output
            grad = np.broadcast_to(grad, x.shape)
        else:
            grad = grad_output * np.ones_like(x)
        
        return grad * scale


class Reshape(Function):
    """改变形状"""
    def forward(self, x, *shape):
        self.save_for_backward(x.shape)
        return x.reshape(*shape)
    
    def backward(self, grad_output):
        original_shape = self.saved_tensors[0]
        return grad_output.reshape(original_shape)


class Transpose(Function):
    """转置操作"""
    def forward(self, x, axes):
        self.save_for_backward(axes)
        return np.transpose(x, axes)
    
    def backward(self, grad_output):
        axes = self.saved_tensors[0]
        # 计算逆转置：反向应用相同的转置
        if axes is None:
            inv_axes = None
        else:
            inv_axes = [0] * len(axes)
            for i, a in enumerate(axes):
                inv_axes[a] = i
        return np.transpose(grad_output, inv_axes)


class Log(Function):
    """自然对数，添加小量防止数值问题"""
    def forward(self, x):
        self.save_for_backward(x)
        return np.log(x + 1e-8)
    
    def backward(self, grad_output):
        x = self.saved_tensors[0]
        return grad_output / (x + 1e-8)


class Exp(Function):
    """指数函数"""
    def forward(self, x):
        self.save_for_backward(x)
        return np.exp(x)
    
    def backward(self, grad_output):
        x = self.saved_tensors[0]
        return grad_output * np.exp(x)


class Negative(Function):
    """取负"""
    def forward(self, x):
        return -x
    
    def backward(self, grad_output):
        return -grad_output


class Power(Function):
    """幂运算"""
    def forward(self, x, power):
        self.save_for_backward(x, power)
        return x ** power
    
    def backward(self, grad_output):
        x, power = self.saved_tensors
        return grad_output * power * (x ** (power - 1))


class Div(Function):
    """除法运算"""
    def forward(self, x, y):
        self.save_for_backward(x, y)
        return x / y
    
    def backward(self, grad_output):
        x, y = self.saved_tensors
        grad_x = grad_output / y
        grad_y = -grad_output * x / (y ** 2)
        return _reduce_grad_to_shape(grad_x, x.shape), _reduce_grad_to_shape(grad_y, y.shape)


# ============================================================================
# 运算符重载：让Tensor支持 + - * / @ 等操作符
# ============================================================================

def _add(self, other):
    if not isinstance(other, Tensor):
        other = Tensor(other)
    return Add().apply(self, other)

def _mul(self, other):
    if not isinstance(other, Tensor):
        other = Tensor(other)
    return Mul().apply(self, other)

def _matmul(self, other):
    if not isinstance(other, Tensor):
        other = Tensor(other)
    return MatMul().apply(self, other)

def _neg(self):
    return Negative().apply(self)

def _sub(self, other):
    return self + (-other)

def _truediv(self, other):
    if not isinstance(other, Tensor):
        other = Tensor(other)
    return Div().apply(self, other)

def _pow(self, power):
    return Power().apply(self, power)

# 将运算符绑定到Tensor类
Tensor.__add__ = _add
Tensor.__mul__ = _mul
Tensor.__matmul__ = _matmul
Tensor.__neg__ = _neg
Tensor.__sub__ = _sub
Tensor.__truediv__ = _truediv
Tensor.relu = lambda self: ReLU().apply(self)
Tensor.sum = lambda self, axis=None, keepdims=False: Sum().apply(self, axis, keepdims)
Tensor.mean = lambda self, axis=None, keepdims=False: Mean().apply(self, axis, keepdims)
Tensor.log = lambda self: Log().apply(self)
Tensor.exp = lambda self: Exp().apply(self)


# ============================================================================
# 第三部分：卷积层（使用 im2col 实现）
# ============================================================================

def im2col(img, kernel_h, kernel_w, stride=1, pad=0):
    """
    将图像转换为列矩阵，用于高效实现卷积
    
    原理：将每个卷积窗口展平成一列，所有列组成一个大矩阵
    这样卷积运算就变成了矩阵乘法
    
    参数:
        img: 输入图像 [N, C, H, W] (N:批量大小, C:通道数, H:高度, W:宽度)
        kernel_h, kernel_w: 卷积核尺寸
        stride: 步长
        pad: 填充
    
    返回:
        cols: [N*out_h*out_w, C*kernel_h*kernel_w]
              每一行对应一个卷积窗口的展平向量
    """
    # 应用填充
    if pad > 0:
        img = np.pad(img, ((0,0), (0,0), (pad,pad), (pad,pad)), mode='constant')
    
    N, C, H, W = img.shape
    # 计算输出特征图的尺寸
    out_h = (H - kernel_h) // stride + 1
    out_w = (W - kernel_w) // stride + 1
    
    # 初始化列矩阵
    cols = np.zeros((N, out_h, out_w, C * kernel_h * kernel_w))
    
    # 对每个位置提取窗口并展平
    for n in range(N):
        for i in range(out_h):
            for j in range(out_w):
                h_start = i * stride
                w_start = j * stride
                patch = img[n, :, h_start:h_start+kernel_h, w_start:w_start+kernel_w]
                cols[n, i, j, :] = patch.flatten()
    
    return cols.reshape(-1, C * kernel_h * kernel_w)


def col2im(cols, img_shape, kernel_h, kernel_w, stride=1, pad=0):
    """
    将列矩阵转换回图像（用于反向传播）
    
    参数:
        cols: 列矩阵 [N*out_h*out_w, C*kernel_h*kernel_w]
        img_shape: 原始图像形状 (N, C, H, W)
        kernel_h, kernel_w: 卷积核尺寸
        stride: 步长
        pad: 填充
    """
    N, C, H, W = img_shape
    H_pad = H + 2 * pad
    W_pad = W + 2 * pad
    out_h = (H_pad - kernel_h) // stride + 1
    out_w = (W_pad - kernel_w) // stride + 1
    
    # 初始化梯度图像（带填充）
    img_grad = np.zeros((N, C, H_pad, W_pad))
    cols = cols.reshape(N, out_h, out_w, -1)
    
    # 将每个列累加到对应的图像区域
    for n in range(N):
        for i in range(out_h):
            for j in range(out_w):
                h_start = i * stride
                w_start = j * stride
                patch = cols[n, i, j, :].reshape(C, kernel_h, kernel_w)
                img_grad[n, :, h_start:h_start+kernel_h, w_start:w_start+kernel_w] += patch
    
    # 移除填充
    if pad > 0:
        img_grad = img_grad[:, :, pad:-pad, pad:-pad]
    
    return img_grad


class Conv2d:
    """
    二维卷积层
    
    使用im2col技术将卷积转换为矩阵乘法，提高效率
    """
    
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        self.stride = stride
        self.padding = padding
        
        # Xavier初始化：保持梯度方差一致
        # 公式：Uniform(-limit, limit)，其中 limit = sqrt(6/(fan_in+fan_out))
        limit = np.sqrt(6.0 / (in_channels * self.kernel_size[0] * self.kernel_size[1] + out_channels))
        self.weight = Tensor(
            np.random.uniform(-limit, limit, 
                            (out_channels, in_channels, self.kernel_size[0], self.kernel_size[1])),
            requires_grad=True
        )
        self.bias = Tensor(np.zeros(out_channels), requires_grad=True)
        self.parameters = [self.weight, self.bias]
    
    def __call__(self, x):
        return Conv2dFunction().apply(self, x)


class Conv2dFunction(Function):
    """卷积运算的函数包装器，实现前向和反向传播"""
    
    def forward(self, layer, x):
        """
        前向传播：使用im2col + 矩阵乘法实现卷积
        """
        self.layer = layer
        self.save_for_backward(x)
        
        N, C, H, W = x.shape
        k_h, k_w = layer.kernel_size
        stride = layer.stride
        padding = layer.padding
        
        # 步骤1：将输入图像转换为列矩阵
        x_cols = im2col(x.data, k_h, k_w, stride, padding)
        
        # 步骤2：将卷积核reshape为2D矩阵 [out_channels, in_channels*k_h*k_w]
        w = layer.weight.data.reshape(layer.out_channels, -1)
        
        # 步骤3：矩阵乘法实现卷积
        out_h = (H + 2*padding - k_h) // stride + 1
        out_w = (W + 2*padding - k_w) // stride + 1
        out = w @ x_cols.T  # [out_channels, N*out_h*out_w]
        
        # 步骤4：添加偏置
        out = out + layer.bias.data.reshape(-1, 1)
        
        # 步骤5：重塑为输出形状 [N, out_channels, out_h, out_w]
        out = out.reshape(layer.out_channels, N, out_h, out_w).transpose(1, 0, 2, 3)
        
        # 保存用于反向传播
        self.x_shape = x.shape
        self.x_cols = x_cols
        
        return out
    
    def backward(self, grad_output):
        """
        反向传播：计算权重、偏置和输入的梯度
        """
        layer = self.layer
        N, C, H, W = self.x_shape
        k_h, k_w = layer.kernel_size
        stride = layer.stride
        padding = layer.padding
        
        # 重塑输出梯度 [N, out_channels, out_h, out_w] -> [out_channels, N*out_h*out_w]
        grad_output_reshaped = grad_output.transpose(1, 0, 2, 3).reshape(layer.out_channels, -1)
        
        # 权重的梯度：dL/dW = grad_output @ x_cols
        grad_weight = grad_output_reshaped @ self.x_cols
        grad_weight = grad_weight.reshape(layer.out_channels, C, k_h, k_w)
        
        # 偏置的梯度：对输出梯度求和
        grad_bias = np.sum(grad_output_reshaped, axis=1)
        
        # 输入的梯度：dL/dx = W.T @ grad_output，然后col2im转换
        w_reshaped = layer.weight.data.reshape(layer.out_channels, -1)
        grad_x_cols = w_reshaped.T @ grad_output_reshaped
        grad_x = col2im(grad_x_cols, (N, C, H, W), k_h, k_w, stride, padding)
        
        # 更新参数的梯度
        if layer.weight.requires_grad:
            if layer.weight.grad is None:
                layer.weight.grad = grad_weight
            else:
                layer.weight.grad += grad_weight
        
        if layer.bias.requires_grad:
            if layer.bias.grad is None:
                layer.bias.grad = grad_bias
            else:
                layer.bias.grad += grad_bias
        
        return None, grad_x  # 第一个返回值为layer的梯度（None，因为layer不是Tensor）


# ============================================================================
# 第四部分：神经网络层
# ============================================================================

class Linear:
    """
    全连接层（线性层）
    实现 y = x @ W + b
    """
    
    def __init__(self, in_features, out_features):
        # Xavier初始化
        limit = np.sqrt(6.0 / (in_features + out_features))
        self.weight = Tensor(
            np.random.uniform(-limit, limit, (in_features, out_features)),
            requires_grad=True
        )
        self.bias = Tensor(np.zeros(out_features), requires_grad=True)
        self.parameters = [self.weight, self.bias]
    
    def __call__(self, x):
        # 确保输入是二维的
        if x.data.ndim == 1:
            x = x.reshape(-1, x.shape[0])
        result = x @ self.weight + self.bias
        return result


class Dropout:
    """
    Dropout层
    
    训练时以概率p随机将神经元输出置0，并缩放1/(1-p)保持期望不变
    测试时不进行dropout，只使用完整的网络
    
    作用：防止过拟合，提高模型泛化能力
    """
    
    def __init__(self, p=0.5):
        self.p = p          # 失活概率
        self.training = True
        self.mask = None    # 保存掩码，反向传播时不需要梯度
    
    def __call__(self, x):
        if not self.training or self.p == 0:
            return x
        
        # 生成掩码：以概率1-p保留，并缩放
        # 缩放因子1/(1-p)保持期望不变：E[mask*x] = (1-p)*x/(1-p) = x
        self.mask = (np.random.rand(*x.shape) > self.p) / (1 - self.p)
        return x * self.mask
    
    def train(self):
        self.training = True
    
    def eval(self):
        self.training = False


class BatchNorm:
    """
    批量归一化层
    
    对每个特征维度进行归一化：使数据分布稳定，加速训练
    
    公式：y = gamma * (x - mean) / sqrt(var + eps) + beta
    
    训练时使用batch的均值和方差，同时更新running_mean/running_var
    测试时使用累积的running_mean和running_var
    """
    
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        self.num_features = num_features
        self.eps = eps              # 防止除零的小量
        self.momentum = momentum    # 运行时统计的更新动量
        
        # 可学习参数：缩放和平移
        self.gamma = Tensor(np.ones(num_features), requires_grad=True)
        self.beta = Tensor(np.zeros(num_features), requires_grad=True)
        
        # 运行时统计（用于推理）
        self.running_mean = np.zeros(num_features)
        self.running_var = np.ones(num_features)
        
        self.training = True
        self.parameters = [self.gamma, self.beta]
    
    def __call__(self, x):
        return BatchNormFunction().apply(self, x)


class BatchNormFunction(Function):
    """BatchNorm的函数包装器，实现前向和反向传播"""
    
    def forward(self, layer, x):
        self.layer = layer
        
        # 处理不同维度的输入
        # 对于卷积层：[N, C, H, W] -> reshape为 [C, N*H*W]
        if x.ndim > 2:
            N, C, H, W = x.shape
            x_reshaped = x.transpose(1, 0, 2, 3).reshape(C, -1)
        else:
            # 对于全连接层：[N, C] -> [C, N]
            N = x.shape[0]
            C = x.shape[1] if x.ndim > 1 else 1
            x_reshaped = x.T if x.ndim > 1 else x.reshape(1, -1)
        
        self.original_shape = x.shape
        self.x_ndim = x.ndim
        self.N = N
        self.C = C
        
        if layer.training:
            # 计算当前batch的均值和方差
            mean = np.mean(x_reshaped, axis=1, keepdims=True)
            var = np.var(x_reshaped, axis=1, keepdims=True)
            
            # 更新运行时统计（使用指数移动平均）
            layer.running_mean = (1 - layer.momentum) * layer.running_mean + layer.momentum * mean.flatten()
            layer.running_var = (1 - layer.momentum) * layer.running_var + layer.momentum * var.flatten()
        else:
            # 推理时使用累积的统计量
            mean = layer.running_mean.reshape(-1, 1)
            var = layer.running_var.reshape(-1, 1)
        
        # 归一化
        x_norm = (x_reshaped - mean) / np.sqrt(var + layer.eps)
        
        # 缩放和平移
        gamma = layer.gamma.data.reshape(-1, 1)
        beta = layer.beta.data.reshape(-1, 1)
        out = gamma * x_norm + beta
        
        # 保存用于反向传播
        self.x_reshaped = x_reshaped
        self.mean = mean
        self.var = var
        self.std_inv = 1.0 / np.sqrt(var + layer.eps)
        
        # 恢复原始形状
        if x.ndim > 2:
            out = out.reshape(C, N, H, W).transpose(1, 0, 2, 3)
        elif x.ndim == 2:
            out = out.T
        else:
            out = out.flatten()
        
        return out
    
    def backward(self, grad_output):
        """
        BatchNorm反向传播的推导较复杂
        需要计算输入、gamma、beta的梯度
        """
        layer = self.layer
        
        # 重塑梯度到 [C, N*...] 形状
        if self.x_ndim > 2:
            _, C, H, W = self.original_shape
            grad_reshaped = grad_output.transpose(1, 0, 2, 3).reshape(self.C, -1)
        elif self.x_ndim == 2:
            grad_reshaped = grad_output.T
        else:
            grad_reshaped = grad_output.reshape(1, -1)
        
        N = self.N
        
        # 中心化后的输入
        x_centered = self.x_reshaped - self.mean
        
        # 计算输入梯度
        # 公式推导较复杂，最终结果：
        # grad_norm = grad_output * gamma
        # grad_var = sum(grad_norm * x_centered * (-0.5 * std_inv^3))
        # grad_mean = sum(grad_norm * (-std_inv)) + grad_var * mean(-2*x_centered)
        # grad_x = grad_norm * std_inv + grad_var * (2*x_centered/N) + grad_mean/N
        
        grad_norm = grad_reshaped * layer.gamma.data.reshape(-1, 1)
        grad_var = np.sum(grad_norm * x_centered * -0.5 * self.std_inv**3, axis=1, keepdims=True)
        grad_mean = np.sum(grad_norm * -self.std_inv, axis=1, keepdims=True) + grad_var * np.mean(-2.0 * x_centered, axis=1, keepdims=True)
        
        grad_x = grad_norm * self.std_inv + grad_var * (2.0 * x_centered / N) + grad_mean / N
        
        # gamma和beta的梯度
        grad_gamma = np.sum(grad_reshaped * (self.x_reshaped - self.mean) * self.std_inv, axis=1)
        grad_beta = np.sum(grad_reshaped, axis=1)
        
        # 恢复输入梯度的形状
        if self.x_ndim > 2:
            grad_x = grad_x.reshape(self.C, N, H, W).transpose(1, 0, 2, 3)
        elif self.x_ndim == 2:
            grad_x = grad_x.T
        else:
            grad_x = grad_x.flatten()
        
        # 更新参数梯度
        if layer.gamma.requires_grad:
            if layer.gamma.grad is None:
                layer.gamma.grad = grad_gamma
            else:
                layer.gamma.grad += grad_gamma
        
        if layer.beta.requires_grad:
            if layer.beta.grad is None:
                layer.beta.grad = grad_beta
            else:
                layer.beta.grad += grad_beta
        
        return None, grad_x


class Sequential:
    """
    顺序容器，将多个层按顺序组合
    
    类似于PyTorch的nn.Sequential
    """
    
    def __init__(self, *layers):
        self.layers = layers
        # 收集所有层的参数
        self.parameters = []
        for layer in layers:
            if hasattr(layer, 'parameters'):
                self.parameters.extend(layer.parameters)
    
    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x
    
    def train(self):
        """设置为训练模式"""
        for layer in self.layers:
            if hasattr(layer, 'training'):
                layer.training = True
    
    def eval(self):
        """设置为评估模式"""
        for layer in self.layers:
            if hasattr(layer, 'training'):
                layer.training = False


# ============================================================================
# 第五部分：损失函数
# ============================================================================

class CrossEntropyLoss:
    """
    交叉熵损失函数（用于分类任务）
    
    组合了Softmax + 负对数似然
    Loss = -log(softmax(logits)[正确类别])
    """
    
    def __call__(self, logits, targets):
        """
        参数:
            logits: [N, C] 未经过softmax的原始分数
            targets: [N] 真实类别索引（0到C-1）
        """
        if not isinstance(targets, Tensor):
            targets = Tensor(targets)
        
        return CrossEntropyFunction().apply(logits, targets.data)


class CrossEntropyFunction(Function):
    """交叉熵损失的函数包装器"""
    
    def forward(self, logits, targets):
        # 数值稳定的Softmax：减去最大值防止溢出
        logits_shifted = logits - np.max(logits, axis=1, keepdims=True)
        exp_logits = np.exp(logits_shifted)
        probs = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)
        
        # 计算损失
        N = logits.shape[0]
        batch_indices = np.arange(N)
        target_indices = targets.astype(np.int32)
        correct_probs = probs[batch_indices, target_indices]
        
        loss = -np.log(correct_probs + 1e-8)
        loss_value = np.mean(loss)
        
        # 保存用于反向传播
        self.save_for_backward(probs, targets, N)
        
        return np.array(loss_value, dtype=np.float32)
    
    def backward(self, grad_output):
        """
        交叉熵 + Softmax 的梯度公式：
        dL/dlogits = probs - one_hot(target)
        """
        probs, targets, N = self.saved_tensors
        
        grad = probs.copy()
        batch_indices = np.arange(N)
        target_indices = targets.astype(np.int32)
        grad[batch_indices, target_indices] -= 1
        grad = grad / N * grad_output
        
        return grad, None  # logits的梯度，targets的梯度为None


class MSELossFunction(Function):
    """均方误差损失的函数包装器"""
    
    def forward(self, pred, target):
        self.save_for_backward(pred, target)
        diff = pred - target
        loss = np.mean(diff ** 2)
        return np.array(loss, dtype=np.float32)
    
    def backward(self, grad_output):
        pred, target = self.saved_tensors
        if pred.shape != target.shape:
            target = target.reshape(pred.shape)
        N = pred.shape[0]
        diff = pred - target
        grad = 2 * diff / N * grad_output
        return grad, None


class MSELoss:
    """均方误差损失"""
    
    def __call__(self, pred, target):
        if not isinstance(target, Tensor):
            target = Tensor(target)
        return MSELossFunction().apply(pred, target)


# ============================================================================
# 第六部分：优化器
# ============================================================================

class SGD:
    """
    随机梯度下降优化器
    
    支持动量和权重衰减
    """
    
    def __init__(self, parameters, lr=0.01, momentum=0, weight_decay=0):
        self.parameters = parameters
        self.lr = lr                      # 学习率
        self.momentum = momentum          # 动量系数
        self.weight_decay = weight_decay  # 权重衰减系数（L2正则化）
        self.velocities = [np.zeros_like(p.data) for p in parameters]  # 动量缓存
    
    def zero_grad(self):
        """清零所有参数的梯度"""
        for param in self.parameters:
            param.grad = None
    
    def step(self):
        """更新参数"""
        for i, param in enumerate(self.parameters):
            if param.grad is not None:
                # 添加权重衰减（L2正则化）
                grad = param.grad + self.weight_decay * param.data
                
                if self.momentum > 0:
                    # 带动量的更新：v = momentum*v - lr*grad
                    self.velocities[i] = self.momentum * self.velocities[i] - self.lr * grad
                    param.data += self.velocities[i]
                else:
                    # 标准SGD更新
                    param.data -= self.lr * grad


class Adam:
    """
    Adam优化器（Adaptive Moment Estimation）
    
    结合了动量和自适应学习率的优点
    """
    
    def __init__(self, parameters, lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
        self.parameters = parameters
        self.lr = lr
        self.betas = betas          # (beta1, beta2) 动量衰减系数
        self.eps = eps              # 防止除零
        self.weight_decay = weight_decay
        
        # 一阶动量（均值）和二阶动量（未中心化的方差）
        self.m = [np.zeros_like(p.data) for p in parameters]
        self.v = [np.zeros_like(p.data) for p in parameters]
        self.t = 0                  # 时间步数
    
    def zero_grad(self):
        for param in self.parameters:
            param.grad = None
    
    def step(self):
        self.t += 1
        
        for i, param in enumerate(self.parameters):
            if param.grad is not None:
                # 添加权重衰减
                grad = param.grad + self.weight_decay * param.data
                
                # 更新动量
                self.m[i] = self.betas[0] * self.m[i] + (1 - self.betas[0]) * grad
                self.v[i] = self.betas[1] * self.v[i] + (1 - self.betas[1]) * (grad ** 2)
                
                # 偏差修正（因为初始时为0，早期会偏向0）
                m_hat = self.m[i] / (1 - self.betas[0] ** self.t)
                v_hat = self.v[i] / (1 - self.betas[1] ** self.t)
                
                # 更新参数
                update = self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
                if update.shape != param.data.shape:
                    update = update.reshape(param.data.shape)
                param.data -= update


# ============================================================================
# 第七部分：完整使用示例
# ============================================================================

def example_conv_and_batchnorm():
    """卷积层和BatchNorm示例"""
    print("="*50)
    print("示例1：卷积层 + BatchNorm")
    print("="*50)
    
    # 创建输入 [batch=2, channels=1, height=28, width=28]
    x = Tensor(np.random.randn(2, 1, 28, 28))
    
    # 创建卷积层 + BatchNorm
    conv = Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
    bn = BatchNorm(32)
    
    # 前向传播
    out = conv(x)
    out = bn(out)
    out = out.relu()
    
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {out.shape}")
    print(f"卷积层权重形状: {conv.weight.shape}")
    print(f"BatchNorm gamma形状: {bn.gamma.shape}")
    print(f"BatchNorm beta形状: {bn.beta.shape}")
    
    # 测试反向传播
    loss = out.sum()
    loss.backward()
    print(f"反向传播完成")
    print(f"卷积层权重梯度形状: {conv.weight.grad.shape if conv.weight.grad is not None else 'None'}")
    print(f"BatchNorm gamma梯度形状: {bn.gamma.grad.shape if bn.gamma.grad is not None else 'None'}")
    print(f"BatchNorm beta梯度形状: {bn.beta.grad.shape if bn.beta.grad is not None else 'None'}")


def example_cross_entropy():
    """交叉熵损失示例"""
    print("\n" + "="*50)
    print("示例2：交叉熵损失")
    print("="*50)
    
    # 模拟分类任务：3个样本，5个类别
    logits = Tensor(np.random.randn(3, 5))
    targets = Tensor(np.array([0, 2, 4]))  # 真实类别
    
    loss_fn = CrossEntropyLoss()
    loss = loss_fn(logits, targets)
    
    print(f"Logits形状: {logits.shape}")
    print(f"目标: {targets.data}")
    print(f"损失值: {loss.item():.4f}")
    
    # 反向传播
    loss.backward()
    print(f"Logits梯度形状: {logits.grad.shape if logits.grad is not None else 'None'}")


def example_dropout():
    """Dropout示例"""
    print("\n" + "="*50)
    print("示例3：Dropout层（训练 vs 测试）")
    print("="*50)
    
    x = Tensor(np.ones((1, 10)))
    dropout = Dropout(p=0.5)
    
    # 训练模式
    dropout.training = True
    out_train = dropout(x)
    print(f"训练模式输出: {out_train.data}")
    print(f"训练模式非零比例: {np.mean(out_train.data != 0):.2f}")
    
    # 测试模式
    dropout.training = False
    out_test = dropout(x)
    print(f"测试模式输出: {out_test.data}")


def example_adam():
    """Adam优化器示例"""
    print("\n" + "="*50)
    print("示例4：Adam优化器 vs SGD")
    print("="*50)
    
    # 简单的线性回归
    np.random.seed(42)
    X_data = np.random.randn(100, 2)
    y_data = X_data @ np.array([[2.0], [3.0]]) + 1.0 + np.random.randn(100, 1) * 0.1
    
    X = Tensor(X_data)
    y = Tensor(y_data)
    
    # 使用Adam的模型
    model_adam = Sequential(Linear(2, 1))
    criterion = MSELoss()
    optimizer_adam = Adam(model_adam.parameters, lr=0.1)
    
    # 使用SGD的模型
    model_sgd = Sequential(Linear(2, 1))
    optimizer_sgd = SGD(model_sgd.parameters, lr=0.1)
    
    # 训练
    epochs = 100
    for epoch in range(epochs):
        # Adam
        pred_adam = model_adam(X)
        loss_adam = criterion(pred_adam, y)
        optimizer_adam.zero_grad()
        loss_adam.backward()
        optimizer_adam.step()
        
        # SGD
        pred_sgd = model_sgd(X)
        loss_sgd = criterion(pred_sgd, y)
        optimizer_sgd.zero_grad()
        loss_sgd.backward()
        optimizer_sgd.step()
        
        if epoch % 20 == 0:
            print(f"Epoch {epoch:3d} - Adam Loss: {loss_adam.item():.6f}, SGD Loss: {loss_sgd.item():.6f}")
    
    print(f"\n训练结果 - Adam: 权重={model_adam.layers[0].weight.data.flatten()}, 偏置={model_adam.layers[0].bias.data[0]:.4f}")
    print(f"训练结果 - SGD:  权重={model_sgd.layers[0].weight.data.flatten()}, 偏置={model_sgd.layers[0].bias.data[0]:.4f}")


def example_full_model():
    """完整模型示例：CNN + Dropout + BatchNorm"""
    print("\n" + "="*50)
    print("示例5：完整的卷积神经网络")
    print("="*50)
    
    class SimpleCNN:
        """简单的卷积神经网络用于分类"""
        def __init__(self):
            # 第一个卷积块：1通道 -> 32通道
            self.conv1 = Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
            self.bn1 = BatchNorm(32)
            
            # 第二个卷积块：32通道 -> 64通道
            self.conv2 = Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
            self.bn2 = BatchNorm(64)
            
            # Dropout防止过拟合
            self.dropout = Dropout(0.5)
            
            # 全连接层：64*28*28 -> 10类
            self.fc = Linear(64 * 28 * 28, 10)
            
            # 收集所有参数
            self.parameters = (self.conv1.parameters + self.bn1.parameters + 
                              self.conv2.parameters + self.bn2.parameters + 
                              self.fc.parameters)
        
        def __call__(self, x):
            x = self.conv1(x)
            x = self.bn1(x)
            x = x.relu()
            
            x = self.conv2(x)
            x = self.bn2(x)
            x = x.relu()
            
            x = self.dropout(x)
            x = x.reshape(x.shape[0], -1)  # 展平
            x = self.fc(x)
            return x
        
        def train(self):
            """切换到训练模式"""
            self.dropout.training = True
            self.bn1.training = True
            self.bn2.training = True
        
        def eval(self):
            """切换到评估模式"""
            self.dropout.training = False
            self.bn1.training = False
            self.bn2.training = False
    
    # 创建模型和数据
    model = SimpleCNN()
    x = Tensor(np.random.randn(4, 1, 28, 28))  # batch=4, 1通道, 28x28图像
    targets = Tensor(np.array([0, 1, 2, 3]))
    
    # 训练模式
    model.train()
    logits = model(x)
    loss_fn = CrossEntropyLoss()
    loss = loss_fn(logits, targets)
    
    print(f"训练模式 - 输入形状: {x.shape}")
    print(f"训练模式 - Logits形状: {logits.shape}")
    print(f"训练模式 - 损失: {loss.item():.4f}")
    
    # 反向传播
    optimizer = Adam(model.parameters, lr=0.001)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print("反向传播和参数更新完成！")
    
    # 评估模式
    model.eval()
    logits_eval = model(x)
    print(f"评估模式 - Logits形状: {logits_eval.shape}")


if __name__ == "__main__":
    # 运行所有示例
    example_conv_and_batchnorm()
    example_cross_entropy()
    example_dropout()
    example_adam()
    example_full_model()