%matplotlib inline
import torch
import torchvision
from torch.utils import data
from torchvision import transforms
from d2l import torch as d2l

d2l.use_svg_display()

# 通过ToTensor实例将图像数据从PIL类型变换成32位浮点数格式，
# 并除以255使得所有像素的数值均在0～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)

len(mnist_train), len(mnist_test)

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
    """
    绘制图像列表。
    
    参数:
    imgs: 一个图像集合，可以是PIL图像或者是PyTorch张量。
    num_rows: 要显示的行数。
    num_cols: 要显示的列数。
    titles: （可选）一个标题列表，用于每张图像。
    scale: （可选）图像显示的缩放比例，默认为1.5。
    """
    # 计算整个画布的大小
    figsize = (num_cols * scale, num_rows * scale)
    # 创建一个子图序列
    _, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize)
    # 将axes数组扁平化，方便遍历
    axes = axes.flatten()
    # 遍历图像和对应的轴（axes）进行绘图
    for i, (ax, img) in enumerate(zip(axes, imgs)):
        if torch.is_tensor(img):
            # 如果图像是PyTorch张量，则将其转换为numpy数组并显示
            ax.imshow(img.numpy())
        else:
            # 如果图像是PIL图像，则直接显示
            ax.imshow(img)
        # 隐藏x和y轴的标签
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        # 如果提供了标题，为每张图像设置标题
        if titles:
            ax.set_title(titles[i])
    # 返回绘制的轴（axes），以便于进一步的操作或调整
    return axes

X, y = next(iter(data.DataLoader(mnist_train, batch_size=18)))
show_images(X.reshape(18, 28, 28), 2, 9, titles=get_fashion_mnist_labels(y));

batch_size = 256

def get_dataloader_workers():  #@save
    """使用4个进程来读取数据"""
    return 4

train_iter = data.DataLoader(mnist_train, batch_size, shuffle=True,
                             num_workers=get_dataloader_workers())

timer = d2l.Timer()
for X, y in train_iter:
    continue
f'{timer.stop():.2f} sec'

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()))

import torch
from IPython import display
from d2l import torch as d2l

batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

# 定义输入向量的维度，对于FashionMNIST数据集，每张图像是28x28像素，
# 因此当图像被展平成向量时，它的大小是784
num_inputs = 784

# 定义输出向量的维度，这里有10个输出类别，对应于FashionMNIST数据集的10种服装类别
num_outputs = 10

# 初始化权重矩阵W。
# 使用正态分布随机初始化，均值为0，标准差为0.01。
# W的形状是(num_inputs, num_outputs)，这样每个输入特征都会与每个输出类别关联一个权重。
W = torch.normal(0, 0.01, size=(num_inputs, num_outputs), requires_grad=True)

# 初始化偏置项b。
# 由于每个输出类别都有一个偏置，所以b的形状是(num_outputs,)。
# 初始化为0，并且设置requires_grad=True，以便于在反向传播时计算梯度。
b = torch.zeros(num_outputs, requires_grad=True)

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

def net(X):
    return softmax(torch.matmul(X.reshape((-1, W.shape[0])), W) + b)

y = torch.tensor([0, 2])
y_hat = torch.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
y_hat[[0, 1], y]

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

cross_entropy(y_hat, y)

def accuracy(y_hat, y):  #@save
    """
    计算预测正确的数量。
    
    参数:
    y_hat: 模型的预测输出。可以是每个类别的得分或概率。
    y: 真实的标签。
    
    返回:
    预测正确的样本数占总样本数的比例。
    """
    # 检查预测输出是否为二维且类别数大于1，
    # 如果是，则使用argmax获取每个样本最高得分的类别索引作为预测类别
    if len(y_hat.shape) > 1 and y_hat.shape[1] > 1:
        y_hat = y_hat.argmax(axis=1)
    
    # 比较预测类别（y_hat）和真实标签（y）是否相同，
    # 结果是一个布尔型张量，其中True表示预测正确，False表示预测错误。
    cmp = y_hat.type(y.dtype) == y
    
    # 将布尔型张量转换为与y相同的数据类型（通常是整型），
    # 然后计算True的数量（即预测正确的数量）。
    # 最后，将其转换为float类型，以便计算准确率。
    return float(cmp.type(y.dtype).sum())

def evaluate_accuracy(net, data_iter):  #@save
    """
    计算在指定数据集上模型的精度。
    
    参数:
    net: 要评估的模型，可以是任何具有前向传播功能的PyTorch模型。
    data_iter: 数据迭代器，用于提供评估数据集的批次数据。
    
    返回:
    模型在整个数据集上的准确率。
    """
    # 检查是否为PyTorch的nn.Module模型。如果是，则将模型设置为评估模式。
    # 评估模式会关闭模型中的Dropout和BatchNorm等层的训练行为，以保证前向传播的一致性。
    if isinstance(net, torch.nn.Module):
        net.eval()
    
    # 初始化一个累加器实例用于跟踪正确预测的数量和预测的总数量。
    metric = Accumulator(2)  # 正确预测数、预测总数
    
    # 不计算梯度，以节省计算资源，因为在评估模式下不需要进行反向传播。
    with torch.no_grad():
        for X, y in data_iter:
            # 计算当前批次的准确率，并累加到metric中。
            # net(X)调用模型进行预测，accuracy函数计算准确率。
            metric.add(accuracy(net(X), y), y.numel())
    
    # 计算整个数据集上的总准确率。
    # metric[0]是正确预测的数量，metric[1]是预测的总数量。
    return metric[0] / metric[1]

class Accumulator:  #@save
    """在n个变量上累加"""
    
def __init__(self, n):
    """
    类的初始化方法。
    参数:
    n (int): 指定需要累加的变量的数量。
    属性:
    data (list of float): 用于存储每个变量的累加值的列表，初始时每个变量的值都设置为0.0。
    """
    self.data = [0.0] * n

def add(self, *args):
    """
    将传入的值加到累加器的对应变量上。
  
    参数:
    *args: 可变数量的参数，每个参数对应于需要累加到的变量的增量值。
    
    说明:
    - 使用zip函数将self.data中的当前累加值和传入的增量值配对。
    - 对每对值进行相加操作，并更新self.data中的累加值。
    """
    self.data = [a + float(b) for a, b in zip(self.data, args)]

def reset(self):
    """
    重置累加器，将所有变量的累加值重置为0.0。
    """
    self.data = [0.0] * len(self.data)

def __getitem__(self, idx):
    """
    使累加器支持下标访问操作。
    
    参数:
    idx (int): 需要访问的变量的索引。
    
    返回:
    float: 指定索引处变量的当前累加值。
    """
    return self.data[idx]

def train_epoch_ch3(net, train_iter, loss, updater):  #@save
    """
    训练模型一个迭代周期。
参数:
net: 训练的神经网络模型。
train_iter: 训练数据的迭代器。
loss: 损失函数。
updater: 参数更新器，可以是PyTorch的优化器或自定义的更新函数。

返回:
训练过程中的平均损失和平均准确率。
"""
# 如果net是PyTorch的神经网络模型，设置为训练模式。
# 这对于某些特定层如Dropout和BatchNorm层在训练和评估阶段行为不同是必要的。
if isinstance(net, torch.nn.Module):
    net.train()

# 使用Accumulator实例来跟踪损失和准确率的累积值，以及处理的总样本数。
metric = Accumulator(3)  # 分别对应损失总和、准确度总和、样本数

# 遍历训练数据迭代器
for X, y in train_iter:
    # 前向传播：计算预测值
    y_hat = net(X)
    # 计算损失
    l = loss(y_hat, y)
    
    # 如果updater是PyTorch优化器，则使用优化器进行反向传播和参数更新
    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 train_ch3(net, train_iter, test_iter, loss, num_epochs, updater):  #@save
    """
    训练模型。
    
    参数:
    net: 要训练的神经网络模型。
    train_iter: 训练数据集的迭代器。
    test_iter: 测试数据集的迭代器。
    loss: 使用的损失函数。
    num_epochs: 训练的总周期数。
    updater: 参数更新器，可以是PyTorch的优化器或自定义的更新函数。
    """
    # 初始化动画生成器，用于可视化训练过程。
    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):
        # 训练模型一个迭代周期，并返回训练损失和训练精度。
        train_metrics = train_epoch_ch3(net, train_iter, loss, updater)
        # 在测试集上评估模型的精度。
        test_acc = evaluate_accuracy(net, test_iter)
        # 使用动画生成器添加当前周期的训练损失、训练精度和测试精度。
        animator.add(epoch + 1, train_metrics + (test_acc,))
    
    # 训练结束后，获取最后一个周期的训练损失和训练精度。
    train_loss, train_acc = train_metrics
    
    # 断言语句用于检查训练过程是否达到预期的效果。
    # 检查训练损失是否小于0.5。
    assert train_loss < 0.5, train_loss
    # 检查训练精度是否在合理范围内（大于0.7，小于等于1）。
    assert train_acc <= 1 and train_acc > 0.7, train_acc
    # 检查测试精度是否在合理范围内（大于0.7，小于等于1）。
    assert test_acc <= 1 and test_acc > 0.7, test_acc

lr = 0.1

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

num_epochs = 10
train_ch3(net, train_iter, test_iter, cross_entropy, num_epochs, updater)

def predict_ch3(net, test_iter, n=8):  #@save
    """
    预测标签并显示图像及其真实标签和预测标签。
    
    参数:
    net: 训练好的神经网络模型。
    test_iter: 测试数据集的迭代器。
    n: 要展示的图像数量，默认为8。
    """
    # 从测试数据迭代器中取出一批数据。
    # 这里只取第一批数据用于展示。
    for X, y in test_iter:
        break
    
    # 获取这批数据的真实标签，并转换为对应的文字标签。
    trues = d2l.get_fashion_mnist_labels(y)
    
    # 使用模型对这批数据进行预测，获取预测结果的最大值索引作为预测标签，
    # 并转换为对应的文字标签。
    preds = d2l.get_fashion_mnist_labels(net(X).argmax(axis=1))
    
    # 为每张图像准备标题，包含真实标签和预测标签。
    titles = [true +'\n' + pred for true, pred in zip(trues, preds)]
    
    # 使用d2l.show_images函数显示图像及其标题。
    # 图像被重塑为28x28的大小，只展示前n张图像。
    d2l.show_images(
        X[0:n].reshape((n, 28, 28)), 1, n, titles=titles[0:n])
    
    predict_ch3(net, test_iter)

import torch
from torch import nn
from d2l import torch as d2l

batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

# PyTorch不会隐式地调整输入的形状。因此，
# 我们在线性层前定义了展平层（flatten），来调整网络输入的形状
net = nn.Sequential(nn.Flatten(), nn.Linear(784, 10))

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

net.apply(init_weights);

loss = nn.CrossEntropyLoss(reduction='none')

trainer = torch.optim.SGD(net.parameters(), lr=0.1)

num_epochs = 10
d2l.train_ch3(net, train_iter, test_iter, loss, num_epochs, trainer)