1. 数据读取与可视化

import os
import time
import torch
import torchvision
from torch.utils import data
from torchvision import transforms
import matplotlib.pyplot as plt
from matplotlib_inline import backend_inline


# --- 解决环境问题的配置 ---
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# --- 类和函数的定义 ---
class Timer:
    """一个简单的计时器类"""

    def __init__(self):
        self.start_time = time.time()

    def stop(self):
        return time.time() - self.start_time


def use_svg_display():
    """使用svg格式在Jupyter中显示绘图。"""
    backend_inline.set_matplotlib_formats('svg')


def get_fashion_mnist_labels(labels):
    """将数字标签转换为文本标签。"""
    text_labels = ['T恤', '裤子', '套头衫', '连衣裙', '外套',
                   '凉鞋', '衬衫', '运动鞋', '包', '短靴']
    return [text_labels[int(i)] for i in labels]


def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5):
    """在一个网格中绘制多张图像。"""
    # ... (函数内部代码不变)
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    axes = axes.flatten()
    for i, (ax, img) in enumerate(zip(axes, imgs)):
        if torch.is_tensor(img):
            ax.imshow(img.numpy())
        else:
            ax.imshow(img)
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        if titles:
            ax.set_title(titles[i])
    return axes


def get_dataloader_workers():
    """子进程数量"""
    return 0


def load_data_fashion_mnist(batch_size, resize=None):  #@save
    """下载Fashion-MNIST数据集，然后将其加载到内存中"""
    trans = [transforms.ToTensor()]
    if resize:
        trans.insert(0, transforms.Resize(resize))
    trans = transforms.Compose(trans)
    mnist_train = torchvision.datasets.FashionMNIST(
        root="../data", train=True, transform=trans, download=True)
    mnist_test = torchvision.datasets.FashionMNIST(
        root="../data", train=False, transform=trans, download=True)
    return (data.DataLoader(mnist_train, batch_size, shuffle=True,
                            num_workers=get_dataloader_workers()),
            data.DataLoader(mnist_test, batch_size, shuffle=False,
                            num_workers=get_dataloader_workers()))

if __name__ == '__main__':

    # 1. 加载数据集
    trans = transforms.ToTensor()
    mnist_train = torchvision.datasets.FashionMNIST(
        root="../data", train=True, transform=trans, download=True)
    mnist_test = torchvision.datasets.FashionMNIST(
        root="../data", train=False, transform=trans, download=True)

    # 2. 打印数据集信息 (现在它只会被主进程打印一次)
    print(f"训练集样本数: {len(mnist_train)}")
    print(f"测试集样本数: {len(mnist_test)}")
    print(f"单张图片的张量形状: {mnist_train[0][0].shape}")

    # 3. 可视化一小批图像
    batch_size_vis = 18
    vis_iter = data.DataLoader(mnist_train, batch_size=batch_size_vis, shuffle=True)
    X_vis, y_vis = next(iter(vis_iter))
    show_images(X_vis.reshape(18, 28, 28), 2, 9, titles=get_fashion_mnist_labels(y_vis))
    plt.show()

    # 4. 测试数据加载速度
    batch_size_test = 256
    num_workers = get_dataloader_workers()
    train_iter = data.DataLoader(mnist_train, batch_size_test, shuffle=True,
                                 num_workers=num_workers)

    print("\n--- 开始测试数据加载速度 ---")
    timer = Timer()
    for X, y in train_iter:
        continue

    print(f'使用 {num_workers} 个进程加载 {len(mnist_train)} 张图片总共花费: {timer.stop():.2f} 秒')

• trans = transforms.ToTensor(): 创建一个 ToTensor 转换器的实例。它的作用是：

  1. 将输入的 PIL 图像（Python Imaging Library 格式）或 NumPy 数组转换为 PyTorch 张量 (torch.Tensor)。
  2. 将图像的像素值从 [0, 255] 的整数范围，自动缩放到 [0.0, 1.0] 的浮点数范围。这是通过将每个像素值除以 255 实现的。数据归一化是神经网络训练中一个常见的预处理步骤。

• mnist_train = torchvision.datasets.FashionMNIST(...): 加载 Fashion-MNIST 的训练数据集
  1. root="../data": 指定数据集存放的目录。如果该目录下没有数据集，PyTorch 会自动下载。
  2. train=True: 表示我们要加载的是训练集。
  3. transform=trans: 告诉数据集加载器，对每一张读入的图片都要应用我们前面定义的 trans 转换，也就是将它变成归一化后的张量。
  4. download=True: 如果在 root 目录下找不到数据集文件，就自动从网上下载。

• mnist_test = torchvision.datasets.FashionMNIST(...): 以同样的方式加载测试数据集，唯一的区别是 train=False。

运行代码得到

训练集样本数: 60000
测试集样本数: 10000
单张图片的张量形状: torch.Size([1, 28, 28])
— 开始测试数据加载速度 —
使用 4 个进程加载 60000 张图片总共花费: 7.48 秒

2. softmax回归的从零开始实现

# softmax_complex.py
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
import torch
import time
from IPython import display
from torchvision import transforms
import torchvision
from torch.utils import data
import matplotlib.pyplot as plt
from matplotlib_inline import backend_inline
from IPython import display

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def get_dataloader_workers():
    """子进程数量"""
    return 0


def load_data_fashion_mnist(batch_size, resize=None):  # @save
    """下载Fashion-MNIST数据集，然后将其加载到内存中"""
    trans = [transforms.ToTensor()]
    if resize:
        trans.insert(0, transforms.Resize(resize))
    trans = transforms.Compose(trans)
    mnist_train = torchvision.datasets.FashionMNIST(
        root="../data", train=True, transform=trans, download=True)
    mnist_test = torchvision.datasets.FashionMNIST(
        root="../data", train=False, transform=trans, download=True)
    return (data.DataLoader(mnist_train, batch_size, shuffle=True,
                            num_workers=get_dataloader_workers()),
            data.DataLoader(mnist_test, batch_size, shuffle=False,
                            num_workers=get_dataloader_workers()))


def softmax(X):
    X_exp = torch.exp(X)
    partition = X_exp.sum(1, keepdim=True)
    return X_exp / partition  # 这里应用了广播机制


class Accumulator:  # @save
    """在n个变量上累加"""

    def __init__(self, n):
        self.data = [0.0] * n

    def add(self, *args):
        self.data = [a + float(b) for a, b in zip(self.data, args)]

    def reset(self):
        self.data = [0.0] * len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def cross_entropy(y_hat, y):
    return - torch.log(y_hat[range(len(y_hat)), y])


def accuracy(y_hat, y):  # @save
    """计算预测正确的数量"""
    if len(y_hat.shape) > 1 and y_hat.shape[1] > 1:
        y_hat = y_hat.argmax(axis=1)
    cmp = y_hat.type(y.dtype) == y
    return float(cmp.type(y.dtype).sum())


def evaluate_accuracy(net, data_iter, device=None):
    """计算模型精度 (修改版，支持device)"""
    if isinstance(net, torch.nn.Module):
        net.eval()
    metric = Accumulator(2)
    with torch.no_grad():
        for X, y in data_iter:
            if device:
                X, y = X.to(device), y.to(device)
            metric.add(accuracy(net(X), y), y.numel())
    return metric[0] / metric[1]


class Animator:
    """在动画中绘制数据。"""

    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None,
                 ylim=None, xscale='linear', yscale='linear',
                 fmts=('-', 'm--', 'g-.', 'r:'), nrows=1, ncols=1,
                 figsize=(3.5, 2.5)):

        # 1. 替换 d2l.use_svg_display()
        def use_svg_display():
            """使用svg格式在Jupyter中显示绘图。"""
            backend_inline.set_matplotlib_formats('svg')

        use_svg_display()

        # 2. 替换 d2l.plt.subplots()
        self.fig, self.axes = plt.subplots(nrows, ncols, figsize=figsize)
        if nrows * ncols == 1:
            # 如果只有一个子图，plt.subplots返回的不是列表，我们把它统一成列表
            self.axes = [self.axes, ]

        # 3. 替换 d2l.set_axes()
        #    我们将 set_axes 的逻辑直接封装在 config_axes 方法中
        self.xlabel = xlabel
        self.ylabel = ylabel
        self.legend = legend if legend is not None else []
        self.xlim = xlim
        self.ylim = ylim
        self.xscale = xscale
        self.yscale = yscale
        self.fmts = fmts

        # 定义一个内部方法来配置坐标轴
        self._config_axes()

        # 初始化数据存储
        self.X, self.Y = None, None

    def _config_axes(self):
        """配置坐标轴。"""
        # 我们只配置第一个子图，因为 d2l 的 Animator 也是这样做的
        ax = self.axes[0]
        if self.xlabel:
            ax.set_xlabel(self.xlabel)
        if self.ylabel:
            ax.set_ylabel(self.ylabel)
        if self.xlim:
            ax.set_xlim(self.xlim)
        if self.ylim:
            ax.set_ylim(self.ylim)
        ax.set_xscale(self.xscale)
        ax.set_yscale(self.yscale)
        if self.legend:
            ax.legend(self.legend)
        ax.grid()

    def add(self, x, y):
        """向图表中添加多个数据点。"""
        # 确保 y 是一个列表或元组
        if not hasattr(y, "__len__"):
            y = [y]
        n = len(y)
        # 确保 x 是一个列表或元组，且长度与 y 相同
        if not hasattr(x, "__len__"):
            x = [x] * n

        # 初始化数据存储列表
        if not self.X:
            self.X = [[] for _ in range(n)]
        if not self.Y:
            self.Y = [[] for _ in range(n)]

        # 添加新的数据点
        for i, (a, b) in enumerate(zip(x, y)):
            if a is not None and b is not None:
                self.X[i].append(a)
                self.Y[i].append(b)

        # --- 绘图部分 ---
        ax = self.axes[0]
        # 1. 清除当前子图上的旧内容
        ax.cla()
        # 2. 绘制所有线条
        for x_data, y_data, fmt in zip(self.X, self.Y, self.fmts):
            ax.plot(x_data, y_data, fmt)

        # 3. 重新配置坐标轴的标签、范围等
        self._config_axes()

        # 4. 在Jupyter中显示和刷新图像
        display.display(self.fig)
        display.clear_output(wait=True)


def train_epoch_ch3(net, train_iter, loss, updater, device=None):
    """训练模型一个迭代周期 (修改版，支持device)"""
    if isinstance(net, torch.nn.Module):
        net.train()
    metric = Accumulator(3)
    for X, y in train_iter:
        # 将数据移动到指定的设备上
        if device:
            X, y = X.to(device), y.to(device)

        y_hat = net(X)
        l = loss(y_hat, y)
        if isinstance(updater, torch.optim.Optimizer):
            updater.zero_grad()
            l.mean().backward()
            updater.step()
        else:
            l.sum().backward()
            updater(X.shape[0])
        metric.add(float(l.sum()), accuracy(y_hat, y), y.numel())
    return metric[0] / metric[2], metric[1] / metric[2]


def sgd(params, lr, batch_size):
    """手动实现的小批量随机梯度下降（原地操作）。"""
    with torch.no_grad():
        for param in params:
            # 原地更新参数
            param -= lr * param.grad / batch_size
            # 原地清空梯度
            param.grad.zero_()


def train_ch3(net, train_iter, test_iter, loss, num_epochs, updater, device=None):
    """训练模型 (修改版，支持device)"""
    animator = Animator(xlabel='epoch', xlim=[1, num_epochs], ylim=[0.3, 0.9],
                        legend=['train loss', 'train acc', 'test acc'])
    for epoch in range(num_epochs):
        # 将 device 参数传递下去
        train_metrics = train_epoch_ch3(net, train_iter, loss, updater, device)
        # evaluate_accuracy 也需要 device 参数
        test_acc = evaluate_accuracy(net, test_iter, device)
        animator.add(epoch + 1, train_metrics + (test_acc,))
    train_loss, train_acc = train_metrics
    assert train_loss < 0.5, train_loss
    assert train_acc <= 1 and train_acc > 0.7, train_acc
    assert test_acc <= 1 and test_acc > 0.7, test_acc


def run(device=None):
    """
    封装了完整的训练和测试流程，并指定在哪个设备上运行。
    :param device: torch.device, 指定计算设备 (e.g., torch.device('cuda'))
    """
    # 1. 设置超参数
    lr = 0.1
    batch_size = 512
    num_epochs = 10

    # 2. 加载数据
    train_iter, test_iter = load_data_fashion_mnist(batch_size)

    # 3. 初始化模型参数
    num_inputs = 784
    num_outputs = 10

    # 将 W 和 b 直接创建在指定的设备上
    W = torch.normal(0, 0.01, size=(num_inputs, num_outputs), requires_grad=True, device=device)
    b = torch.zeros(num_outputs, requires_grad=True, device=device)

    # 4. 定义模型和优化器
    #    这个 net 函数现在会捕获在 run 函数作用域内的 W 和 b
    def net(X):
        # 注意：这里的 X 也需要被移动到和 W, b 相同的设备上
        # 我们将在 train_epoch_ch3 中处理
        return softmax(torch.matmul(X.reshape((-1, W.shape[0])), W) + b)

    def updater(batch_size):
        return sgd([W, b], lr, batch_size)

    # 5. 开始训练并计时
    print(f"--- 开始在 {device} 上训练 ---")
    start_time = time.time()
    train_ch3(net, train_iter, test_iter, cross_entropy, num_epochs, updater, device)
    end_time = time.time()

    print(f"--- 在 {device} 上的训练结束 ---")
    print(f"总耗时: {end_time - start_time:.2f} 秒\n")
    # 清除 Animator 的图像，为下一次运行做准备


if __name__ == '__main__':
    # (此处省略了您已有的所有类和函数的定义代码)
    # 确保您的脚本里包含了所有必要的定义

    # 运行 CPU 版本
    run(device=torch.device('cpu'))

    # 检查 CUDA 是否可用，如果可用，则运行 GPU 版本
    if torch.cuda.is_available():
        run(device=torch.device('cuda'))
    else:
        print("CUDA (GPU) 在此设备上不可用。")
    plt.show()

因为数据量比较小，所以用cpu与gpu没什么区别，主要的时间还是花在了数据加载上

— 在 cpu 上的训练结束 —
总耗时: 58.27 秒
— 在 cuda 上的训练结束 —
总耗时: 67.53 秒

3. softmax回归调包实现

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
import softmax_complex
from torch import nn
import torch
import time
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False



def init_weights(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, std=0.01)


batch_size = 256
train_iter, test_iter = softmax_complex.load_data_fashion_mnist(batch_size)
net = nn.Sequential(nn.Flatten(), nn.Linear(784, 10))
net.apply(init_weights)
loss = nn.CrossEntropyLoss(reduction='none')
trainer = torch.optim.SGD(net.parameters(), lr=0.1)
num_epochs = 10
start_time = time.time()
softmax_complex.train_ch3(net, train_iter, test_iter, loss, num_epochs, trainer)
end_time = time.time()
print(f"总耗时: {end_time - start_time:.2f} 秒\n")
plt.show()

总耗时: 60.89 秒