my_train.py

import os
import time
import torch
import numpy as np
import torchvision
from torchvision import transforms
import matplotlib.pyplot as plt
import multiprocessing  # 用于获取CPU核心数

# import mplcursors

torch.set_num_threads(min(8, max(multiprocessing.cpu_count() - 2, 1)))


def try_gpu(i=0):  # @save
    """如果存在，则返回dml设备，否则返回cpu"""
    try:
        import torch_directml

        if torch_directml.device_count() > i:
            return torch_directml.device(i)
    except:
        return torch.device("cpu")


def set_seed(seed=42):
    if seed is None:
        seed = int(time.time() * 1000 + os.getpid()) % (2**32 - 1)
    torch.manual_seed(seed)
    np.random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


device = try_gpu()
# scaler = torch.GradScaler(device=device.type)


def get_dataloader_workers():
    return min(8, multiprocessing.cpu_count())


def load_data_fashion_mnist(batch_size, resize=None):
    """下载Fashion-MNIST数据集, 然后将其加载到内存中"""
    trans = [transforms.ToTensor()]
    if resize:
        trans.insert(0, transforms.Resize(resize))  # type: ignore
    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
    )
    use_pin_memory = device != torch.device("cpu")
    return (
        torch.utils.data.DataLoader(
            mnist_train,
            batch_size,
            shuffle=True,
            num_workers=get_dataloader_workers(),
            pin_memory=use_pin_memory,
            persistent_workers=True,
        ),
        torch.utils.data.DataLoader(
            mnist_test,
            batch_size,
            shuffle=False,
            num_workers=get_dataloader_workers(),
            pin_memory=use_pin_memory,
            persistent_workers=True,
        ),
    )


# 设置 matplotlib 使用 SVG 格式（替代 d2l.use_svg_display）
def use_svg_display(enable_color=True):
    """完整替代 d2l.use_svg_display()"""
    plt.rcParams["figure.figsize"] = (3.5, 2.5)
    plt.rcParams["figure.dpi"] = 120
    plt.rcParams["figure.autolayout"] = True
    plt.rcParams["savefig.format"] = "svg"  # 保存为SVG
    plt.rcParams["image.interpolation"] = "nearest"  # 最近邻插值保持锐利
    plt.rcParams["axes.grid"] = True  # 禁用网格线
    plt.rcParams["image.cmap"] = "gray"  # 使用灰度图显示

    if enable_color:
        # 彩色模式配置
        plt.rcParams["image.cmap"] = "viridis"  # 翠绿色更优的默认彩色映射
        plt.rcParams["figure.facecolor"] = "white"  # 白色背景
        plt.rcParams["axes.facecolor"] = "white"  # 白色坐标背景
    else:
        # 灰度模式配置
        plt.rcParams["image.cmap"] = "gray"
        plt.rcParams["figure.facecolor"] = "white"
        plt.rcParams["axes.facecolor"] = "white"

    plt.rcParams["font.family"] = "DejaVu Sans"  # 更好支持字符
    plt.rcParams["svg.fonttype"] = "none"  # 确保SVG中文字可编辑
    # plt.rcParams["image.composite_image"] = False  # 确保精确颜色渲染


use_svg_display(enable_color=True)


def cross_entropy(y_hat, y):  # 交叉熵
    return -torch.log(y_hat[range(len(y_hat)), y])  # 高级索引


def accuracy(y_hat: torch.Tensor, y: torch.Tensor):  # @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((y_hat == y).sum().item())


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 evaluate_accuracy(net, data_iter):  # @save
    """计算在指定数据集上模型的精度"""
    if isinstance(net, torch.nn.Module):
        net.eval()  # 将模型设置为评估模式
    metric = Accumulator(2)  # 正确预测数、预测总数
    with torch.no_grad():
        for X, y in data_iter:
            X, y = X.to(device, non_blocking=True).half(), y.to(
                device, non_blocking=True
            )
            metric.add(accuracy(net(X), y), y.numel())
    return metric[0] / metric[1]


def train_epoch_ch6(net, train_iter, loss, updater, scheduler=None):  # @save
    """训练模型一个迭代周期（定义见第6章）"""
    # 将模型设置为训练模式
    if isinstance(net, torch.nn.Module):
        net.train()
    # 训练损失总和、训练准确度总和、样本数
    metric = Accumulator(3)
    for X, y in train_iter:
        # 计算梯度并更新参数
        X, y = X.to(device, non_blocking=True).half(), y.to(device, non_blocking=True)
        # with torch.autocast(device_type=device.type, dtype=torch.float16):
        y_hat = net(X)
        l = loss(y_hat, y)
        if isinstance(updater, torch.optim.Optimizer):
            # use_amp = scaler is not None
            # if use_amp:
            #     updater.zero_grad()
            #     # 混合精度反向传播和优化
            #     scaler.scale(l).backward()
            #     # # 梯度裁剪（可选）
            #     # scaler.unscale_(updater)
            #     # torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
            #     scaler.step(updater)
            #     scaler.update()
            # else:
            # 使用PyTorch内置的优化器和损失函数
            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())
    # 在每个epoch结束后更新学习率
    if scheduler is not None and isinstance(updater, torch.optim.Optimizer):
        scheduler.step()
    # 返回训练损失和训练精度
    return (metric[0] / metric[2], metric[1] / metric[2]), metric[2]


from IPython import display


class Animator:  # @save
    """在动画中绘制数据"""

    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),
    ):
        # 增量地绘制多条线
        if legend is None:
            legend = []
        use_svg_display()
        self.fig, self.axes = plt.subplots(nrows, ncols, figsize=figsize)
        if nrows * ncols == 1:
            self.axes = [self.axes]

        # self.config_axes = lambda: d2l.set_axes(
        #     self.axes[0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend
        # )
        def create_config_axes(ax):
            def config_func():
                """配置坐标轴属性的闭包函数"""
                if xlabel:
                    ax.set_xlabel(xlabel)
                if ylabel:
                    ax.set_ylabel(ylabel)
                if xlim:
                    ax.set_xlim(xlim)
                if ylim:
                    ax.set_ylim(ylim)
                ax.set_xscale(xscale)
                ax.set_yscale(yscale)
                if legend:
                    ax.legend(legend)

            return config_func

        # try:
        #     self.cursor = mplcursors.cursor(self.fig)
        #     self.cursor.connect("add", self._on_hover)
        # except ImportError:
        #     pass
        self.config_axes = [create_config_axes(ax) for ax in self.axes]
        self.X, self.Y, self.fmts = None, None, fmts

    def add(self, x, y):
        # 向图表中添加多个数据点
        if not hasattr(y, "__len__"):
            y = [y]
        n = len(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)
        self.axes[0].cla()
        for x, y, fmt in zip(self.X, self.Y, self.fmts):
            line = self.axes[0].plot(x, y, fmt)[0]
            color = line.get_color()
            self.axes[0].text(
                x[-1],
                y[-1],
                f"{y[-1]:.2%}",
                color=color,
                bbox=dict(facecolor="white", alpha=0.6),
            )
        self.config_axes[0]()
        display.display(self.fig)
        display.clear_output(wait=True)

    def _on_hover(self, sel):
        """鼠标悬停时显示数据点详情"""
        artist = sel.artist
        idx = sel.target.index
        y_val = artist.get_ydata()[idx]
        sel.annotation.set_text(f"Y: {y_val:.2%}")
        sel.annotation.get_bbox_patch().set(facecolor="yellow", alpha=0.9)


def mytrain_ch6(
    net,
    train_iter,
    test_iter,
    loss,
    num_epochs,
    updater,
    scheduler=None,
    save_model=False,
    use_fp16=True,
):  # @save
    """训练模型（定义见第6章）"""
    print(f"Train on {device}")
    if use_fp16:
        net = net.half()
        # scaler = torch.cuda.amp.GradScaler(enabled=use_fp16)
    animator = Animator(
        xlabel="epoch",
        xlim=[1, num_epochs],
        ylim=[0.2, 1.0],
        legend=["train loss", "train acc", "test acc"],
    )
    start_time = time.time()
    for epoch in range(num_epochs):
        train_metrics, total_metrics = train_epoch_ch6(
            net, train_iter, loss, updater, scheduler
        )
        if test_iter is not None:
            test_acc = evaluate_accuracy(net, test_iter)
            animator.add(epoch + 1, train_metrics + (test_acc,))
        else:
            animator.add(epoch + 1, train_metrics)
    train_loss, train_acc = train_metrics
    total_time = time.time() - start_time
    print(
        f"Spend {total_time}s, {total_metrics * num_epochs / total_time:.1f} examples/sec on {str(device)}"
    )
    # 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
    if save_model:
        torch.save(net, "mytrain.pth")
        print("模型已完整保存到 mytrain.pth")


# 预测
def get_fashion_mnist_labels(labels):  # @save
    """返回Fashion-MNIST数据集的文本标签"""
    text_labels = [
        "t-shirt",
        "trouser",
        "pullover",
        "dress",
        "coat",
        "sandal",
        "shirt",
        "sneaker",
        "bag",
        "ankle boot",
    ]
    return [text_labels[int(i)] for i in labels]


def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5):  # @save
    """绘制图像列表"""
    _fig, axes = plt.subplots(
        num_rows, num_cols, figsize=(num_cols * scale, num_rows * scale)
    )
    axes = axes.flatten()  # 安全展平
    for i, (ax, img) in enumerate(zip(axes, imgs)):
        if torch.is_tensor(img):
            if img.device.type == "privateuseone":  # DirectML设备
                img = img.cpu()
            elif img.device.type != "cpu":
                img = img.detach().cpu()
            img = img.numpy()
        ax.imshow(img)  # PIL图片
        ax.axes.get_xaxis().set_visible(False)  # 隐藏坐标轴
        ax.axes.get_yaxis().set_visible(False)
        if titles:  # 设置标题
            ax.set_title(titles[i])
    return axes


def mypredict_ch6(net, test_iter, n=30, resize=28):  # @save
    """预测标签（定义见第6章）"""
    for X, y in test_iter:
        break
    X = X.to(device).half()
    with torch.no_grad():
        preds_y = net(X).argmax(axis=1).cpu()
    trues = get_fashion_mnist_labels(y)
    preds = get_fashion_mnist_labels(preds_y)
    titles = [true + "\n" + pred for true, pred in zip(trues, preds)]
    show_images(X[0:n].reshape((n, resize, resize)), 3, n // 3, titles=titles[0:n])


def load_full_model(path="full_model.pth"):
    """导入模型"""
    model = torch.load(path, weights_only=False)
    model.eval()
    return model


# Net
from torch import nn
from torch.nn import functional as F


class Residual(nn.Module):  # @save
    def __init__(self, input_channels, num_channels, use_1x1conv=False, strides=1):
        super().__init__()
        self.conv1 = nn.Conv2d(
            input_channels, num_channels, kernel_size=3, padding=1, stride=strides
        )
        self.conv2 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
        if use_1x1conv:
            self.conv3 = nn.Conv2d(
                input_channels, num_channels, kernel_size=1, stride=strides
            )
        else:
            self.conv3 = None
        self.bn1 = nn.BatchNorm2d(num_channels)
        self.bn2 = nn.BatchNorm2d(num_channels)

    def forward(self, X):
        Y = F.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3:
            X = self.conv3(X)
        Y += X
        return F.relu(Y)


def my_resnet18(num_classes, in_channels):
    """改进的ResNet-18模型"""

    def resnet_block(input_channels, num_channels, num_residuals, first_block=False):
        blk = []
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.append(
                    Residual(input_channels, num_channels, use_1x1conv=True, strides=2)
                )
            else:
                blk.append(Residual(num_channels, num_channels))
        return blk

    b1 = nn.Sequential(
        nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
        nn.BatchNorm2d(64),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
    )
    b2 = nn.Sequential(*resnet_block(64, 64, 2, first_block=True))
    b3 = nn.Sequential(*resnet_block(64, 128, 2))
    b4 = nn.Sequential(*resnet_block(128, 256, 2))
    b5 = nn.Sequential(*resnet_block(256, 512, 2))
    net = nn.Sequential(
        b1,
        b2,
        b3,
        b4,
        b5,
        nn.AdaptiveAvgPool2d((1, 1)),
        nn.Flatten(),
        nn.Linear(512, num_classes),
    )
    return net


def get_net(num_classes, in_channels, module="resnet18"):
    net = my_resnet18(num_classes, in_channels).to(device)
    return net