深度学习实战（基于pytroch）系列（八）softmax回归基于pytorch的代码实现

上一节我们已经实现softmax回归从零开始使用python代码实现，这节我们用pytorch框架来实现，由于pytorch已经实现了大部分功能，所以写起来代码会十分简洁。

导包

这一步和上一节softmax回归从零开始使用python代码实现基本一致

import torch
from torch import nn
from torch import optim
import torchvision
import torchvision.transforms as transforms

数据加载

这一步和上一节softmax回归从零开始使用python代码实现基本一致

batch_size = 256
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))
])train_dataset = torchvision.datasets.FashionMNIST(root='../data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.FashionMNIST(root='../data', train=False, download=True, transform=transform)train_iter = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_iter = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

定义模型

softmax回归的输出层是一个全连接层。因此，我们添加一个输出个数为10的全连接层。我们使用均值为0、标准差为0.01的正态分布随机初始化模型的权重参数。使用Sequential来构建

net = nn.Sequential(nn.Flatten(),nn.Linear(28*28, 10)
)

nn.Flatten() 工作机制

默认（start_dim=1）

import torch
import torch.nn as nn
test_cases = [torch.randn(4),           # 1D: (4,)torch.randn(4, 3),        # 2D: (4, 3)  torch.randn(4, 3, 5),     # 3D: (4, 3, 5)torch.randn(4, 3, 5, 2),  # 4D: (4, 3, 5, 2)
]flatten = nn.Flatten()print("=== nn.Flatten() 默认行为 (start_dim=1) ===")
for i, tensor in enumerate(test_cases):input_shape = tensor.shapeoutput_shape = flatten(tensor).shapedim_change = len(input_shape) - len(output_shape)print(f"输入: {input_shape} → 输出: {output_shape} | 维度变化: {dim_change}")

默认情况下，相当于batch_size不参与展平，可以通过下面输出结果，了解其工作机制：

输入: torch.Size([4]) → 输出: torch.Size([4]) | 维度变化: 0
输入: torch.Size([4,3]) → 输出: torch.Size([4, 3]) | 维度变化: 0
输入: torch.Size([4, 3, 5]) →输出: torch.Size([4, 15]) | 维度变化: -1
输入: torch.Size([4, 3, 5, 2]) → 输出:torch.Size([4, 30]) | 维度变化: -2

初始化权重

使用均值为0、标准差为0.01的正态分布初始化权重

nn.init.normal_(net[1].weight, mean=0, std=0.01)
nn.init.constant_(net[1].bias, val=0)

损失函数和优化器

定义损失函数（PyTorch的CrossEntropyLoss已经包含了softmax）

loss = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.1)

训练模型

num_epochs = 5for epoch in range(num_epochs):train_l_sum, train_acc_sum, n = 0.0, 0.0, 0for X, y in train_iter:# 前向传播y_hat = net(X)l = loss(y_hat, y)# 反向传播optimizer.zero_grad()l.backward()optimizer.step()train_l_sum += l.item()train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()n += y.shape[0]# 测试准确率test_acc = 0.0test_n = 0with torch.no_grad():for X, y in test_iter:test_acc += (net(X).argmax(dim=1) == y).sum().item()test_n += y.shape[0]test_acc /= test_nprint(f'epoch {epoch + 1}, loss {train_l_sum / len(train_iter):.4f}, 'f'train acc {train_acc_sum / n:.3f}, test acc {test_acc:.3f}')