diff --git a/test.py b/test.py new file mode 100644 index 0000000..8ca155c --- /dev/null +++ b/test.py @@ -0,0 +1,95 @@ +import matplotlib +import torch +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.datasets import MNIST +import matplotlib.pyplot as plt +matplotlib.use('TkAgg') + +# 神经网络类 主体 +class Net(torch.nn.Module): + def __init__(self): + super().__init__() + # 四个全连接层 + self.fc1 = torch.nn.Linear(28 * 28, 64) # 输入28*28像素的图像 64个节点 + self.fc2 = torch.nn.Linear(64, 64) + self.fc3 = torch.nn.Linear(64, 64) + self.fc4 = torch.nn.Linear(64, 10) # 输出十个数字类别 + + def forward(self, x): # x:图像输入 + # 每层首先做全连接线性计算 self.fc1 再套上激活函数torch.nn.functional.relu + x = torch.nn.functional.relu(self.fc1(x)) + x = torch.nn.functional.relu(self.fc2(x)) + x = torch.nn.functional.relu(self.fc3(x)) + # 输出层通过softmax进行归一化 + x = torch.nn.functional.log_softmax(self.fc4(x), dim=1) + return x + + +# 下载MNIST数据集 导入数据 +def get_data_loader(is_train): + # tensor 多维数组 张量 + to_tensor = transforms.Compose([transforms.ToTensor()]) + """ + 第一个参数:''标识数据集下载位置 空表示下载到当前目录 + is_train:指定导入训练集 + """ + data_set = MNIST('', is_train, transform=to_tensor, download=True) + # batch_size:一个批次包含15张图片 shuffle=True:数据随机打乱 + return DataLoader(data_set, batch_size=15, shuffle=True) # 返回数据加载器DataLoader() + + +# 评估神经网络识别正确率 +def evaluate(test_data, net): + # 正确预测数量 + n_correct = 0 + # 预测总数 + n_total = 0 + with torch.no_grad(): + # 从测试集依次取出数据 + for (x, y) in test_data: + # 计算神经网络预测值 + outputs = net.forward(x.view(-1, 28 * 28)) + # 对批次结果进行比较 + for i, output in enumerate(outputs): + # 累加正确预测的数量 argmax 计算数列中最大结果,即预测的手写数字结果 + if torch.argmax(output) == y[i]: + n_correct += 1 + n_total += 1 + # 返回预测正确率 + return n_correct / n_total + + +def main(): + # 导入训练集 + train_data = get_data_loader(is_train=True) + # 导入测试集 + test_data = get_data_loader(is_train=False) + # 初始化神经网络 + net = Net() + # 打印最开始的时候的正确率 约为0.1左右 + print('初始神经网络正确率:', evaluate(test_data, net)) + optimizer = torch.optim.Adam(net.parameters(), lr=0.001) + # epoch:轮次 反复使用train_data进行训练神经网络,提高训练集利用率 + for epoch in range(2): + for (x, y) in train_data: + net.zero_grad() # 初始化 + output = net.forward(x.view(-1, 28 * 28)) # 正向传播 + # nll_lose 对数损失函数 与前的softmax对数计算进行对应 + loss = torch.nn.functional.nll_loss(output, y) # 计算差值 + loss.backward() # 反向误差传播 + optimizer.step() # 优化网络参数 + print('第', epoch+1, '次训练后准确率:', evaluate(test_data, net)) + # 随机抽取三张图象显示网络预测结果 + for (n, (x, _)) in enumerate(test_data): + if n > 2: + break + predict = torch.argmax(net.forward(x[0].view(-1, 28 * 28))) + plt.figure(n) + plt.imshow(x[0].view(28, 28)) + plt.title('prediction:' + str(int(predict))) + plt.show() + + +if __name__ == '__main__': + main()