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)

#定义模型
net = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(),
                    nn.Linear(256, 10))

#定义损失函数——交叉熵损失函数
loss = nn.CrossEntropyLoss()

#训练器
trainer = torch.optim.SGD(net.parameters(), lr=0.1)


#训练模型
num_epochs = 3
for epoch in range(num_epochs):
    for X, y in train_iter:
        #在CrossEntropyLoss()y会自动被整理成one-hot形式：net(X).shape=(256*10)，而y.shape=(256)
        l = loss(net(X),y)  #在这里loss会自动对l求sum，使其变成一个标量
        trainer.zero_grad()
        l.backward() 
        trainer.step()

在这里的关键是net = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(),
                    nn.Linear(256, 10))

nn.Flatten()——展平层，使得数据集的能够输入到模型中去

nn.Linear(784, 256), nn.ReLU()——这两个加在一起叫做hidden layer

nn.Linear(256, 10)——输出层

在这里关键是要知道每层的shape为什么是这样表示的——可以用矩阵的方式去尝试还原这个过程