1.问题描述
利用MLP对Fashion-MNIST数据集进行分类。
2.导入模块
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import sys
from tqdm import tqdm
from matplotlib import pyplot as plt
3. 加载数据集
def load_dataset(batch_size, resize=None):
"""加载Fashion-MNIST数据集"""
transform = transforms.ToTensor()
if resize:
transform.insert(0, transforms.Resize(resize))
transforms.Compose([transform])
mnist_train = datasets.FashionMNIST(root='..\dataset\FashionMnisit',
train=True,
download=False,
transform=transform)
mnist_test = datasets.FashionMNIST(root='..\dataset\FashionMnisit',
train=False,
download=False,
transform=transform)
return (DataLoader(mnist_train, batch_size, shuffle=True, num_workers=4),
DataLoader(mnist_test, batch_size, shuffle=False, num_workers=4))
4.建立网络模型
class MyNet(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 128)
self.fc3 = nn.Linear(128, 64)
self.fc4 = nn.Linear(64, 10)
# 构造Dropout方法,在每次训练过程中都随机“掐死”百分之二十的神经元,防止过拟合。
self.dropout = nn.Dropout(p=0.2)
def forward(self, x):
# 确保输入的tensor是展开的单列数据,把每张图片的通道、长度、宽度三个维度都压缩为一列
x = x.view(x.shape[0], -1)
# 在训练过程中对隐含层神经元的正向推断使用Dropout方法
x = self.dropout(F.relu(self.fc1(x)))
x = self.dropout(F.relu(self.fc2(x)))
x = self.dropout(F.relu(self.fc3(x)))
# 在输出单元不需要使用Dropout方法
x = self.fc4(x)
return x
5.定义训练函数
def train(model, optimizer, data_loader, epoch):
model.train()
loss_function = torch.nn.CrossEntropyLoss()
accu_loss = torch.zeros(1) # 累计损失
accu_num = torch.zeros(1) # 累计预测正确的样本数
optimizer.zero_grad()
sample_num = 0
data_loader = tqdm(data_loader, file=sys.stdout)
for step, data in enumerate(data_loader):
images, labels = data
sample_num += images.shape[0]
pred = model(images)
pred_classes = torch.max(pred, dim=1)[1]
accu_num += torch.eq(pred_classes, labels).sum() # 对两个张量的元素进行比较
loss = loss_function(pred, labels)
loss.backward()
accu_loss += loss.detach()
data_loader.desc = "[train epoch {}] loss: {:.6f}, acc: {:.6f}".format(epoch+1,
accu_loss.item() / (step + 1),
accu_num.item() / sample_num)
if not torch.isfinite(loss):
print('WARNING: non-finite loss, ending training ', loss)
sys.exit(1)
optimizer.step()
optimizer.zero_grad()
return accu_loss.item() / (step + 1), accu_num.item() / sample_num
6.定义评估函数
def evaluate(model, data_loader, epoch):
loss_function = torch.nn.CrossEntropyLoss()
model.eval()
accu_num = torch.zeros(1) # 累计预测正确的样本数
accu_loss = torch.zeros(1) # 累计损失
sample_num = 0
data_loader = tqdm(data_loader, file=sys.stdout)
for step, data in enumerate(data_loader):
images, labels = data
sample_num += images.shape[0]
pred = model(images)
pred_classes = torch.max(pred, dim=1)[1]
accu_num += torch.eq(pred_classes, labels).sum() # 对两个张量的元素进行比较
loss = loss_function(pred, labels)
accu_loss += loss
data_loader.desc = "[test epoch {}] loss: {:.6f}, acc: {:.6f}".format(epoch+1,
accu_loss.item() / (step + 1),
accu_num.item() / sample_num)
return accu_loss.item() / (step + 1), accu_num.item() / sample_num
7.主函数
if __name__ == '__main__':
# 数据集实例化
train_iter, test_iter = load_dataset(batch_size=256)
# 网络实例化
net = MyNet()
# 优化器
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
#
train_losses, test_losses, trainAccuracy, testAccuracy = [], [], [], []
print("------------开始训练-----------")
for epoch in range(15):
# train
train_loss, train_acc = train(model=net,
optimizer=optimizer,
data_loader=train_iter,
epoch=epoch)
# test
test_loss, test_acc = evaluate(model=net,
data_loader=test_iter,
epoch=epoch)
train_losses.append(train_loss)
test_losses.append(test_loss)
trainAccuracy.append(train_acc)
testAccuracy.append(test_acc)
# 保存最后一次迭代的模型
torch.save(net.state_dict(), "./model-{}.pth".format(epoch + 1))
## 绘制误差曲线
plt.figure(1)
plt.plot(train_losses, label='Train loss')
plt.plot(test_losses, label='Test loss')
plt.plot(trainAccuracy, label='Train Accuracy')
plt.plot(testAccuracy, label='Test Accuracy')
plt.legend()
plt.show()
## 进行测试
net.eval()
dataiter = iter(test_iter)
images, labels = dataiter.next()
img = images[0]
img = img.reshape((28, 28)).numpy()
plt.figure(2)
plt.imshow(img)
plt.show()
# 将测试图片转为一维的列向量
img = torch.from_numpy(img)
img = img.view(1, 784)
# 进行正向推断,预测图片所在的类别
with torch.no_grad():
output = torch.squeeze(net(img))
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print(output,predict,predict_cla)
labellist = ['T恤', '裤子', '套衫', '裙子', '外套', '凉鞋', '汗衫', '运动鞋', '包包', '靴子']
prediction = labellist[predict_cla]
probability = predict[predict_cla].numpy()
print(f'神经网络猜测图片里是 {prediction},概率为{probability * 100}%')
8.结果展示
9.参考资料
【1】动手学深度学习 Release 2.0.0-beta0.
【2】https://github.com/TommyZihao/zihaopytorch
【3】https://github.com/WZMIAOMIAO/deep-learning-for-image-processing
10.备注
仅供个人学习记录,侵权删。