使用pytorch加载自己的数据集并搭建LeNet5网络进行训练

之前在学习pytorch入门知识的时候拿了师兄的一个数据集来练手，这篇文章记录一下训练的全过程。

数据集格式

使用pytorch加载数据集首先要清楚数据集的格式。我拿到的数据集是读取一个电子显示器上的数字的图片，这个显示器一次显示5个数字，事先已经对图片进行的预处理，将5个数字切割成单个的数字，切割的做法有利于简化问题，直接识别单个数字即可，只需要搭建一个简单的网络即可，就跟解决经典的mnist数字识别问题一样。不同的是加载数据集的方式有些差别。下图为数据集的图片样例：

说一下图片的尺寸，图片的尺寸是24×32，而mnist数据集的图片尺寸是32×32.
数据集的文件目录结构如下：

label里的格式如下：

这里label是在切割成5个图片前生成的，所以label图片文件的名字不可用，需要从train和test文件夹里读入图片文件名。

加载数据集

接下来就是加载数据集的细节了。首先定义dataset的类：

class customData(Dataset):
    def __init__(self, img_path, txt_path, data_transforms=None, loader=default_loader):
        self.img_names = []
        self.img_labels = []
        fr = open(txt_path)
        for line in fr.readlines():  # 逐行读取
            line = line.strip()  # 滤除行首行尾空格
            for n in line[-5:]:  # 遍历连续的五个数字
                self.img_labels.append(int(n))  # 将五个数字都添加为标签

        files = os.listdir(img_path)
        files.sort(key=lambda x: int(x.split('.')[0]))
        for fn in files:
            self.img_names.append(os.path.join(img_path, fn))

        self.data_transforms = data_transforms
        self.loader = loader


    def __len__(self):
        return len(self.img_names)

    def __getitem__(self, item):
        img_name = self.img_names[item]
        label = self.img_labels[item]
        img = self.loader(img_name)

        if self.data_transforms is not None:
            try:
                img = self.data_transforms(img)
            except:
                print("Cannot transform image: {}".format(img_name))
        return img, label

# use PIL Image to read image
def default_loader(path):
    try:
        img = Image.open(path)
        # print(path)
        return img
    except:
        print("Cannot read image: {}".format(path))

上面代码的细节主要有两点：

1. 读取label的方式：读取label.txt文件的每一行，而我们只需要每一行里面的最后五个数字，所以我们使用for n in line[-5:]:来切割出行末5个数字并一一遍历这5个数字，将其添加到img_labels列表中。

# 读取label.txt文件并将label添加到列表中
fr = open(txt_path)
        for line in fr.readlines():  # 逐行读取
            line = line.strip()  # 滤除行首行尾空格
            for n in line[-5:]:  # 遍历连续的五个数字
                self.img_labels.append(int(n))  # 将五个数字都添加为标签

2. 读取图片的方式：由于files = os.listdir(img_path)读取的图片路径是随机的，所以需要对其按文件名进行排序，随后再按照排序好的路径将图片路径添加到img_names列表中。

files = os.listdir(img_path)
        files.sort(key=lambda x: int(x.split('.')[0]))
        for fn in files:
            self.img_names.append(os.path.join(img_path, fn))

实现dataloader

这一步设置batch_size=128，epochs=10，在transforms里只设置ToTensor()的选项：

BATCH_SIZE = 128
EPOCHS = 10 # 总共训练批次

train_datasets = customData(img_path="./label_and_photos/train/",
                            txt_path="./label_and_photos/label_train.txt",
                            data_transforms=transforms.ToTensor())

test_datasets = customData(img_path="./label_and_photos/test/",
                           txt_path="./label_and_photos/label_test.txt",
                           data_transforms=transforms.ToTensor())

train_loader = DataLoader(dataset=train_datasets,
                          batch_size=BATCH_SIZE,
                          shuffle=True)

test_loader = DataLoader(dataset=test_datasets,
                         batch_size=BATCH_SIZE,
                         shuffle=True)

然后在装载完成后，选取其中一个批次的数据进行预览：

images, labels = next(iter(train_loader))
img = torchvision.utils.make_grid(images)

img = img.numpy().transpose(1, 2, 0)

std = [0.5, 0.5, 0.5]
mean = [0.5, 0.5, 0.5]

img = img * std + mean
# print(labels)
print([labels[i] for i in range(64)])
cv2.imshow('win',img)
key_pressed=cv2.waitKey(0)

效果如下图：

搭建LeNet5网络

首先给出LeNet的结构图：

上图输入的图片的尺寸是32×32，我的数据集是24×32，需要做相应的计算，比较简单，代码里注释得很详细，这里就不再详细解释了。

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(  # (1, 24, 32)
                in_channels=1,
                out_channels=6,
                kernel_size=5,
                stride=1,
                padding=0
            ),  # ->(6, 20, 28)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # ->(6, 10, 14)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, 3, 1, 0),  # (16, 8, 12)
            nn.ReLU(),
            nn.MaxPool2d(2)              # (16, 4, 6)
        )
        self.out = nn.Sequential(
            nn.Linear(16 * 4 * 6, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10),
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  # 将输出化成一维向量
        output = self.out(x)
        return output

训练和测试

定义训练和测试函数：

def train(epoch):  # 定义每个epoch的训练细节
    model.train()  # 设置为trainning模式
    train_acc = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = Variable(data), Variable(target)  # 把数据转换成Variable
        optimizer.zero_grad()  # 优化器梯度初始化为零
        output = model(data)  # 把数据输入网络并得到输出，即进行前向传播
        loss = loss_function(output, target)  # 交叉熵损失函数
        loss.backward()  # 反向传播梯度
        optimizer.step()  # 结束一次前传+反传之后，更新参数
        pred = torch.max(output, 1)[1]
        train_correct = (pred == target).sum()
        train_acc += train_correct.item()
        if batch_idx % 10 == 0:  # 准备打印相关信息
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test():
    model.eval()  # 设置为test模式
    test_loss = 0  # 初始化测试损失值为0
    correct = 0  # 初始化预测正确的数据个数为0
    for data, target in test_loader:
        data, target = Variable(data), Variable(target)  # 计算前要把变量变成Variable形式，因为这样子才有梯度

        output = model(data)
        test_loss += loss_function(output, target).item()  # sum up batch loss 把所有loss值进行累加
        pred = output.data.max(1, keepdim=True)[1]  # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).cpu().sum()  # 对预测正确的数据个数进行累加

    test_loss /= len(test_loader.dataset)  # 因为把所有loss值进行过累加，所以最后要除以总得数据长度才得平均loss
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

因为这次任务比较简单，网络层数很少，这里我使用的是CPU来训练（最主要是自己的电脑配置太低了），可以使用GPU来训练，代码需要相应改一下，改动并不多。
定义main函数入口：

if __name__ == "__main__":

    # 实例化一个网络
    model = LeNet5()

    # 定义损失函数和优化器
    loss_function = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        model.parameters(),
        lr=0.001,
        momentum=0.9
    )
    for epoch in range(1, EPOCHS + 1):  # 以epoch为单位进行循环
        train(epoch)

    test()

    torch.save(model, 'model.pth')  # 保存模型

最后，这种数字识别的任务比较简单，LeNet5网络结构已经能取得很不错的识别效果了，在测试集上的识别准确率是96.4%。之前跑完的时候忘了画训练过程中的准确率曲线图，现在懒得再跑一遍了，感兴趣的同学可以自己去跑一下。

参考资料

https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter05_CNN/5.5_lenet