AMD ROCm能否打破N卡CUDA深度学习垄断，实现MNIST手写数字识别？

前言 直接介绍了我折腾在【Windows系统 + AMD显卡】的多重 debuff 下尝试炼丹，最终成功编译出 ROCm 版本的 PyTorch，这意味着获得了进入炼丹界的钥匙。 接下来，为了验证我自己编译的 PyTorch 是否真正具备生产力，我们可以写一个经典的 MNIST 手写数字识别 脚本。 MNIST 手写数字识别（28x28 像素）被称为深度学习领域的“Hello World” 😄 这个脚本会包含： 显卡识别。 张量运算（测试 ROCm 内核）。 模型训练循环。 接下来我们分步介绍一下。 神经网络简介 先来科普一下。 神经网络是受人类大脑神经元连接方式启发的人工智能核心计算模型，由输入层、隐藏层和输出层构成的大量 “人工神经元” 相互连接而成。它通过模拟人脑信息处理逻辑，先由输入层接收图像、文本等原始数据，再经隐藏层逐层提取和转换特征，最后通过输出层输出分类、预测等结果。其核心优势在于具备自主学习能力，通过 “前向传播计算预测结果、反向传播调整连接权重” 的循环训练，不断优化模型精度，无需人工预设特征规则。如今，神经网络已广泛应用于生活各领域，从手机人脸解锁、智能语音助手，到医疗影像诊断、自动驾驶和个性化推荐，成为推动人工智能落地的关键技术。 我在网上看到一张很不错的图，可以清晰看到各种神经网络的结构。 我记得刚学 Deep Learning 的时候，看了《Python 神经网络编程》这本书，里面作者教的是用 NumPy 手写了一个神经网络，包括后面同作者的另一本书《PyTorch 生成对抗网络》里用 PyTorch 实现的纯全连接层 MNIST 识别模型，都属于前馈神经网络（Feedforward Neural Network，FNN）（也常被称为多层感知机 MLP）。 而本文不用这种基础的网络，而是改成使用卷积神经网络（Convolutional Neural Network，CNN），更具体地说是 “卷积层 + 全连接层” 组合的基础 CNN 结构，这是 MNIST 手写数字识别任务中最经典、最常用的网络类型。 前馈神经网络（FNN/MLP） 是最基础的神经网络结构，无循环、无反馈，数据仅从输入层→隐藏层→输出层单向 “前馈” 流动； 处理图像时，会先把二维图像（如 28x28 的 MNIST）展平成一维向量（784 个元素），再通过全连接层（每个神经元与上一层所有神经元相连）进行计算； 卷积神经网络（CNN） 是专为网格状数据（图像、语音）设计的前馈网络变体，核心新增了卷积层、池化层，保留了前馈的核心逻辑，但优化了特征提取方式； 处理图像时不直接展平，而是通过卷积核 “滑动扫描” 图像，保留空间结构信息，再结合池化降维，最后才接入全连接层分类。 区别 两者核心差异在于是否保留空间结构和连接方式：FNN 全连接 + 展平，适合非空间数据；CNN 局部连接 + 权值共享，专为图像等空间数据设计。 维度 前馈神经网络（FNN / 纯全连接） 卷积神经网络（CNN） 数据处理方式 直接展平为一维向量，丢失空间信息 保留二维空间结构，卷积核扫描提取特征 连接方式 全连接：每个神经元与上一层所有神经元相连 局部连接 + 权值共享：卷积核仅连接局部区域，同一核参数复用 参数数量（MNIST） 示例 FNN：784×128 + 128×64 + 64×10 = 109,328 示例 CNN：conv1 (1×32×3×3) + conv2 (32×64×3×3) + fc1 (9216×128) + fc2 (128×10) ≈ 1,219,624（看似多，但实际更大图像下 CNN 参数远少于 FNN） 特征提取能力 只能学习全局特征，无法捕捉局部空间关系（如数字 “8” 的上下圈） 分层提取特征：低层（边缘、线条）→高层（笔画、数字轮廓），精准捕捉空间特征 平移 / 变形鲁棒性 差：数字轻微位移就会被识别为不同特征 强：池化 + 卷积让模型对小位移、变形不敏感 适用场景 表格数据、简单分类（非空间数据） 图像、视频、语音等网格 / 序列空间数据 CNN的优势 保留空间信息，特征提取更精准 FNN 把 28x28 的图像展平成 784 维向量后，数字 “5” 的顶部横线和右侧竖线的位置关系被完全打乱；而 CNN 的卷积层会 “看到” 这些局部结构，比如第一层卷积提取边缘，第二层提取笔画组合，最终识别更准确（MNIST 上 FNN 准确率约 97%，简单 CNN 可达 99%+）。