李沐softmax实现（pytroch）

import torch
import torchvision
from torch import nn
from torch.utils import data
import os
from torch.utils.tensorboard import SummaryWriter

mnist_train=torchvision.datasets.FashionMNIST(root="FashionMNIST",train=True,transform=torchvision.transforms.ToTensor(),download=True)
mnist_test=torchvision.datasets.FashionMNIST(root="FashionMNIST",train=False,transform=torchvision.transforms.ToTensor(),download=True)

print(len(mnist_train),len(mnist_test))

#一个批次读取的数量
batch_size=100
train_iter=data.DataLoader(mnist_train,batch_size,shuffle=True)

#定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.module=nn.Sequential(
            nn.Flatten(),
            nn.Linear(784,10)
        )
    def forward(self,x):
        x=self.module(x)
        return x

#实例化训练模型
net=Net()
#选取交叉熵损失函数
loss_fn=nn.CrossEntropyLoss()
loss_fn=loss_fn.cuda()

#优化算法
optimizer=torch.optim.SGD(net.parameters(),lr=0.01)
#迭代轮数
num_epochs=10


writer=SummaryWriter("keshihau")
step=0
#训练
for epoch in range(num_epochs):
    total_loss=0
    for x,y in train_iter:

        output=net(x)
        loss=loss_fn(output,y)
        optimizer.zero_grad()
        loss.sum().backward()  #把每一类的loss加和成标量，对加和后的标量loss求梯度，因为向量求梯度不方便
        optimizer.step()
        #tensorboard 可视化损失
        writer.add_scalar("softmax",loss.sum(),step)
        step+=1
        total_loss=total_loss+loss.sum()
    print("第{0}轮loss={1}".format(epoch+1,total_loss))

writer.close()

构建的Net训练模型的依据

其实就是将读入的数据用nn.Flatten()展平，然后在经过一个线性层nn.Linear()

softmax可视化的loss图像

softmax不懂的可以参考沐神的官网3.4. softmax回归 — 动手学深度学习 2.0.0 documentation (d2l.ai)