pytorch入门二：MNIST手写数字网络的代码封装与可视化

代码地址：https://github.com/DongZhouGu/MNIST-Madon

在学习了简单的卷积神经网络实现MNIST手写数字识别，参考了一些其他的代码和博客，让代码更好的封装下，并增加了可视化的部分。 

目录

1 文件组织架构

2  关于__init__.py

3.数据加载

4 模型定义

5 模型训练

6 模型测试

7 主函数

---

1 文件组织架构

├── data/
│   ├── MNIST
│   └── dataset.py
├── models/
│   ├── __init__.py
│   ├── AlexNet.py
│   ├── MyNet.py
│   └── Net.py
└── utils/
│   ├── __init__.py
│   └── visualize.py
├── main.py
├── requirements.txt
├── README.md

• data/：MNIST数据集存放，数据预处理等
• models/：模型定义，可以有多个模型，例如上面的AlexNet，MyNet，等一个模型对应一个文件
• utils/：可能用到的工具函数，在本次实验中主要是封装了可视化工具
• main.py：主文件，训练和测试程序的入口，可通过argparse来指定不同的参数
• requirements.txt：程序依赖的第三方库
• README.md：提供程序的必要说明

2  关于__init__.py

一个目录如果包含了__init__.py 文件，那么它就变成了一个包（package）。__init__.py可以为空，也可以定义包的属性和方法，但其必须存在，其它程序才能从这个目录中导入相应的模块或函数。例如在models/文件夹下有__init__.py，则在main.py 中就可以

from models.MyNet import MyNet

而如果在models/__init__.py中写入

from .MyNet import MyNet

则在main.py中就可以直接写为：

from models import MyNet

 

相比于from models.MyNet import MyNet更加便捷。

3.数据加载

# 加载训练集
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize(mean=(0.5,), std=(0.5,))  # 数据规范化到正态分布
                   ])),
    batch_size=BATCH_SIZE, shuffle=True)  # 指明批量大小，打乱，这是处于后续训练的需要。

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])),
    batch_size=BATCH_SIZE, shuffle=True)

这是根据pytorch自带的接口进行数据加载，我们也可以专门写一个datasets,使得代码封装更有普适性

import torch
from torchvision import datasets, transforms
class Dataset():
    def __init__(self):
        # 注意这是python 2.0的写法
        super(Dataset, self).__init__()

        # Compose用于将多个transfrom组合起来
        # ToTensor()将像素转换为tensor，并做Min-max归一化，即x'=x-min/max-min
        # 相当于将像素从[0,255]转换为[0,1]
        # Normalize()用均值和标准差对图像标准化处理 x'=(x-mean)/std，加速收敛的作用
        # 这里0.131是图片的均值，0.308是方差，通过对原始图片进行计算得出
        # 想偷懒的话可以直接填Normalize([0.5], [0.5])
        # 另外多说一点，因为MNIST数据集图片是灰度图，只有一个通道，因此这里的均值和方差都只有一个值
        # 若是普通的彩色图像，则应该是三个值，比如Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])