import matplotlib.pyplot as plt
from torchvision.transforms import ToTensor
import torch
from torch import nn # 包含构建神经网络的所有模块
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# 训练数据集
training_data = datasets.FashionMNIST(
root="./data", # 存储测试集和训练集的路径
train=True, # 训练集
download=True, # 如果本机没有数据集，就会下载到 root 目录下。
transform=ToTensor() # 对样本数据进行处理，转换为张量数据
)
# 测试数据集
test_data = datasets.FashionMNIST(
root="./data",
train=False,
download=True,
transform=ToTensor()
)
# 标签字典，一个key键对应一个label
labels_map = {
0: "T-Shirt",
1: "Trouser",
2: "Pullover",
3: "Dress",
4: "Coat",
5: "Sandal",
6: "Shirt",
7: "Sneaker",
8: "Bag",
9: "Ankle Boot",
}
# 设置画布大小
# figure = plt.figure(figsize=(8, 8))
# cols, rows = 3, 3
# for i in range(1, cols * rows + 1):
# # 随机生成一个索引
# sample_idx = torch.randint(len(training_data), size=(1,)).item()
# # 获取样本及其对应的标签
# img, label = training_data[sample_idx]
# figure.add_subplot(rows, cols, i)
# # 设置标题
# plt.title(labels_map[label])
# # 不显示坐标轴
# plt.axis("off")
# # 显示灰度图
# plt.imshow(img.squeeze(), cmap="gray")
# plt.show()
# 训练数据加载器; 根据数据集生成一个迭代对象，用于模型的训练
train_dataloader = DataLoader(
# 定义好的数据集
dataset=training_data,
# 设置批量大小
batch_size=128,
# 线程数，默认为0。在Windows下设置大于0的数可能会报错。
num_workers=0,
# 打乱样本的顺序
shuffle=True)
# 测试数据加载器
test_dataloader = DataLoader(
dataset=test_data,
batch_size=128,
shuffle=True)
# 展示图片和标签
# train_features, train_labels = next(iter(train_dataloader))
# print(f"Feature batch shape: {train_features.size()}")
# print(f"Labels batch shape: {train_labels.size()}")
# img = train_features[0].squeeze()
# label = train_labels[0]
# plt.imshow(img, cmap="gray")
# plt.show()
# print(f"Label: {label}")
# 模型定义
class Network(nn.Module):
def __init__(self):
super(Network, self).__init__() # 执行父类中的 init 函数
self.flatten = nn.Flatten() # 将每个大小为28x28的图像转换为784个像素值的连续数组
self.linear_relu_stack = nn.Sequential(
nn.Linear(in_features=28 * 28, out_features=512), # 线性层
nn.ReLU(),
nn.Linear(in_features=512, out_features=512),
nn.ReLU(),
nn.Linear(in_features=512, out_features=10),
)
def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return logits
# 优化模型参数
def train_loop(dataloader, model, loss_func, optimizer, device):
size = len(dataloader.dataset)
for batch, (X, y) in enumerate(dataloader):
X = X.to(device)
y = y.to(device)
# 前向传播，计算预测值
pred = model(X)
# 计算损失
loss = loss_func(pred, y)
# 反向传播，优化参数
optimizer.zero_grad() # 将模型的梯度归 0
loss.backward() # 用来存储每个参数的损失梯度
optimizer.step() # 梯度调整完以后，调整参数
if batch % 100 == 0:
loss, current = loss.item(), batch * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
# 测试模型性能
def test_loop(dataloader, model, loss_fn, device):
size = len(dataloader.dataset)
num_batches = len(dataloader)
test_loss, correct = 0, 0
with torch.no_grad():
for X, y in dataloader:
X = X.to(device)
y = y.to(device)
# 前向传播，计算预测值
pred = model(X)
# 计算损失
test_loss += loss_fn(pred, y).item()
# 计算准确率
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
print(
f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
# input
if __name__ == '__main__':
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
# 定义模型，并将模型移动到设备上
model = Network().to(device)
# 设置超参数
learning_rate = 1e-3
epochs = 20
# 定义损失函数和优化器
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model.parameters(), lr=learning_rate)
# 训练模型
for t in range(epochs):
print(f"Epoch {t + 1}\n-------------------------------")
train_loop(train_dataloader, model, loss_func, optimizer, device)
test_loop(test_dataloader, model, loss_func, device)
print("Done!")
# 保存模型
torch.save(model.state_dict(), 'model_weights.pth')
