我的博客
0%

深度学习(1)自定义图像分类

发表于 分类于

图像分类是深度学习的典型问题，本文使用Mobile-net作为分类网络，实现自定义数据集的分类。主要特点是：

  • 网上教程大多使用构造好的数据集如CIFAR10，缺少如何从头构建自己数据集的过程。本文从采集好的图片开始，一步步构造符合训练要求的数据集。
  • 不熟悉深度学习与计算机视觉的工作者可能对种类繁多的网络结构感到困扰，编写起来更是麻烦。但是Pytorch的torchvision模块已经自带了很多SOTA网络的实现，因此本文直接import了Mobile-net，属于即插即用。
  • 代码讲解较少，如果不熟悉Pytorch请先学习官方图像分类教程，了解基本概念。

综上，读者可以根据本教程快速实现自己的深度学习分类器。

数据准备

分类场景是2020 Running Robot大赛的第一关（上下开横杆），分为两类：一类是横杆关闭，机器人需要等待；一类是横杆打开，机器人可以通过。在机器人头部摄像头的视角下，两类图片如下：

我们在不同位置分别采集这两类数据，构造训练集和验证集，文件目录保持如下（0、1那两个名字可以任取，比如cat dog之类的都无所谓，只要一个文件夹下放一类图片即可）：

接下来我们编写一个preprocess.py文件，构造数据索引文件（这里简便起见使用txt文件）。每一行由“图片路径 + 类别编号”组成，方便构造Dataset类时索引。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import os
import glob
import random

if __name__ == '__main__':
    traindata_path = './data/train'
    labels = os.listdir(traindata_path)

    for index, label in enumerate(labels):
        imglist = glob.glob(os.path.join(traindata_path,label, '*.png'))
        random.shuffle(imglist)

        with open('./data/train.txt', 'a')as f:
            for img in imglist:
                f.write(img + ' ' + str(index))
                f.write('\n')
    

    validdata_path = './data/valid'
    labels = os.listdir(validdata_path)

    for index, label in enumerate(labels):
        imglist = glob.glob(os.path.join(validdata_path,label, '*.png'))
        random.shuffle(imglist)

        with open('./data/valid.txt', 'a')as f:
            for img in imglist:
                f.write(img + ' ' + str(index))
                f.write('\n')

运行完成后，data文件下多了train.txt和valid.txt文件，最好打开检查一下，是否符合“图片路径 + 类别编号”的预期。

自定义Dataset与DataLoader

编写dataset.py文件

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import torch
from torchvision import transforms, datasets
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
from PIL import Image


class  SelfCustomDataset(Dataset):
    def __init__(self, label_file):
        with open(label_file, 'r') as f:
            self.imgs = list(map(lambda line: line.strip().split(' '), f))
        self.transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def __getitem__(self, index):
        img_path, label = self.imgs[index]
        img = Image.open(img_path).convert('RGB')
        img =self.transform(img)
        return img, torch.from_numpy(np.array(int(label)))
 
    def __len__(self):
        return len(self.imgs)


train_datasets = SelfCustomDataset('./data/train.txt')
train_dataloader = torch.utils.data.DataLoader(train_datasets, batch_size=64, shuffle=True, num_workers=1)

val_datasets = SelfCustomDataset('./data/valid.txt')
val_dataloader = torch.utils.data.DataLoader(val_datasets, batch_size=64, shuffle=True, num_workers=1)

这里有两个值得注意的地方：

  • transform里只有两个基本变化，第一个是把Numpy转为Tensor，第二个是把图片归一化。事实上，为了增强鲁棒性，还可以使用更多变换，如剪裁、平移、噪声等。
  • 直接在dataset.py文件中实例化了train_dataloader和val_dataloader，因此在训练代码中from dataset import train_dataloader, val_dataloader即可

训练代码

编写train.py文件

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
import argparse
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models import mobilenet_v2
from dataset import train_dataloader,val_dataloader

save_folder = './model'
os.makedirs(save_folder, exist_ok=True)

model=mobilenet_v2(pretrained=False,num_classes=2)
device = torch.device("cuda:0")
model = model.to(device)

optimizer = optim.Adam(model.parameters(),lr=0.01)
criterion = nn.CrossEntropyLoss()
#criterion = nn.BCEWithLogitsLoss()
# criterion = LabelSmoothSoftmaxCE()
# criterion = LabelSmoothingCrossEntropy()

print("Start Training...")
for epoch in range(20):
    valid_loss = 0.
    for i, data in enumerate(val_dataloader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        valid_loss += loss.item()
    print(valid_loss/(i+1))

    for i, data in enumerate(train_dataloader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

print("Done Training!")

torch.save(model, './model/1.pth')

测试代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import torch
from torchvision import transforms
import numpy as np
import os
from PIL import Image
from dataset import train_dataloader,val_dataloader

device = torch.device("cpu")
model = torch.load('./model/1.pth').to(device)
model.eval()


transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

img = Image.open('data/test/1.png').convert('RGB')
img = transform(img)
img = torch.unsqueeze(img, 0)
with torch.no_grad():
    output = model(img)
pred = np.argmax(output.numpy(),axis=1)
print(pred)

注意： model.eval()非常重要，这会固定batch_norm层参数，否则实际预测效果与训练效果不符。