【阿里云天池】实战：街景符号识别（Baseline）

一、导入包

import os, glob, json
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

import torchvision
from PIL import Image
import numpy as np


import torch
torch.manual_seed(0)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True

import torchvision.models as models
import torchvision.transforms as transforms
import torch.nn as nn
from torch.utils.data.dataset import Dataset

二、数据读取和处理
（1）Dataset

①输入数据：
首先Dataset的输入数据为数据的图像地址、数据的标签和数据是否变换，因为数据的标签是保存在对应的JSON文件里的，它代表的是数据的真实值。

②Init：
保存数据本身的地址，标签，是否变换。

③getitem：
在该方法中，定义了怎么对数据进行处理。

首先使用Image方法按照每一张图片的地址，读取图片信息，并将它保存成RGB通道的形式。

再对图片按照变换参数transform所定义的方法进行图片变换（裁剪，翻转，归一化等）。

最后对突变的标签进行处理，首先将标签变成int形的numpy数据。

lbl = list(lbl) + (5 - len(lbl)) * [10]
首先字符只可能有0-9这十种情况，所以我们可以用数字10代表空字符。同时所有数据中，最长的数据长度为5个字符（只有一个字符是6，可以忽略）。所以通过上述代码可以将一张图片的对应字符可能定长为5个，且空字符为10。
如有一张图的真实信息是19，那么我们在程序中将它变为19XXX，五位定长，而lbl的值就为：“1”,“9”,“10”,“10”,“10”。

class SVHNDataset(Dataset):
    def __init__(self, img_path, img_label, transform=None):
        self.img_path = img_path
        self.img_label = img_label
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None

    def __getitem__(self, index):
        img = Image.open(self.img_path[index]).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        lbl = np.array(self.img_label[index], dtype=np.int)
        lbl = list(lbl) + (5 - len(lbl)) * [10]
        return img, torch.from_numpy(np.array(lbl[:5]))

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

（2）数据的读取和Dataloader
主要是要读取训练数据和测试数据，其中两个数据其实在dataset上没有什么区别。

①数据读取部分
使用glob方法，*.png的方式可以读取文件夹中的所有图片，sort可以给图片排序，使其和JSON数据中的图片信息一一对应。

JSON用来读取JSON数据，并且将JSON数据中的‘label’标签进行一一保存。

train_path = glob.glob('mchar_train/mchar_train/*.png')
train_path.sort()
train_json = json.load(open('mchar_train.json'))
train_label = [train_json[x]['label'] for x in train_json]

②DataLoader部分

train_loader = torch.utils.data.DataLoader(
    SVHNDataset(train_path, train_label,
                transforms.Compose([
                    transforms.Resize((64, 128)),
                    transforms.RandomCrop((60, 120)),
                    transforms.ColorJitter(0.3, 0.3, 0.2),
                    transforms.RandomRotation(10),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])),
        batch_size=10,
    shuffle=True,
    num_workers=10,
)

DataLoader部分，需要输入：对应的地址，标签，变换方法。batch_size，是否打乱，读取线程。

其中变换方法分别为：缩放，随机裁剪，颜色抖动，随机旋转，转为Tensor，归一化。

测试数据的读取差别不大。

val_loader = torch.utils.data.DataLoader(
    SVHNDataset(val_path, val_label,
                transforms.Compose([
                    transforms.Resize((60, 120)),
                    # transforms.ColorJitter(0.3, 0.3, 0.2),
                    # transforms.RandomRotation(5),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])),
    batch_size=10,
    shuffle=False,
    num_workers=10,
)

三、模型
采用RESNET18网络模型，思路是将问题变成定长为5的分类问题，也就是说将进行5次分类。
暂时先不阐述具体内容

class SVHN_Model1(nn.Module):
    def __init__(self):
        super(SVHN_Model1, self).__init__()

        model_conv = models.resnet18(weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1)
        model_conv.avgpool = nn.AdaptiveAvgPool2d(1)
        model_conv = nn.Sequential(*list(model_conv.children())[:-1])
        self.cnn = model_conv

        self.fc1 = nn.Linear(512, 11)
        self.fc2 = nn.Linear(512, 11)
        self.fc3 = nn.Linear(512, 11)
        self.fc4 = nn.Linear(512, 11)
        self.fc5 = nn.Linear(512, 11)

    def forward(self, img):
        feat = self.cnn(img)
        # print(feat.shape)
        feat = feat.view(feat.shape[0], -1)
        c1 = self.fc1(feat)
        c2 = self.fc2(feat)
        c3 = self.fc3(feat)
        c4 = self.fc4(feat)
        c5 = self.fc5(feat)
        return c1, c2, c3, c4, c5

四、训练

def train(train_loader, model, criterion, optimizer, epoch):
    # 切换模型为训练模式
    model.train()
    train_loss = []

    for i, (input, target) in enumerate(train_loader):
        if use_cuda:
            input = input.cuda()
            target = target.cuda()

        c0, c1, c2, c3, c4 = model(input)


        loss = criterion(c0, target[:, 0].long())+ \
               criterion(c1, target[:, 1].long())+ \
               criterion(c2, target[:, 2].long())+ \
               criterion(c3, target[:, 3].long())+ \
               criterion(c4, target[:, 4].long())

        # loss /= 6
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss.append(loss.item())
    return np.mean(train_loss)

五、测试
测试并不需要反馈和更新

def validate(val_loader, model, criterion):
    # 切换模型为预测模型
    model.eval()
    val_loss = []

    # 不记录模型梯度信息
    with torch.no_grad():
        for i, (input, target) in enumerate(val_loader):
            if use_cuda:
                input = input.cuda()
                target = target.cuda()

            c0, c1, c2, c3, c4 = model(input)
            loss = criterion(c0, target[:, 0].long()) + \
                   criterion(c1, target[:, 1].long()) + \
                   criterion(c2, target[:, 2].long()) + \
                   criterion(c3, target[:, 3].long()) + \
                   criterion(c4, target[:, 4].long())
            # loss /= 6
            val_loss.append(loss.item())
    return np.mean(val_loss)

六、预测

def predict(test_loader, model, tta=10):
    model.eval()
    test_pred_tta = None

    # TTA 次数
    for _ in range(tta):
        test_pred = []

        with torch.no_grad():
            for i, (input, target) in enumerate(test_loader):
                if use_cuda:
                    input = input.cuda()

                c0, c1, c2, c3, c4 = model(input)
                if use_cuda:
                    output = np.concatenate([
                        c0.data.cpu().numpy(),
                        c1.data.cpu().numpy(),
                       c2.data.cpu().numpy(),
                        c3.data.cpu().numpy(),
                        c4.data.cpu().numpy()], axis=1)
                else:
                    output = np.concatenate([
                        c0.data.numpy(),
                        c1.data.numpy(),
                        c2.data.numpy(),
                        c3.data.numpy(),
                        c4.data.numpy()], axis=1)

                test_pred.append(output)

        test_pred = np.vstack(test_pred)
        if test_pred_tta is None:
            test_pred_tta = test_pred
        else:
            test_pred_tta += test_pred

    return test_pred_tta

七、主函数

model = SVHN_Model1()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), 0.001)
best_loss = 1000.0

#是否使用GPU
use_cuda = True
model = model.cuda()
if __name__ == '__main__':
    for epoch in range(10):
        train_loss = train(train_loader, model, criterion, optimizer, epoch)
        val_loss = validate(val_loader, model, criterion)

        val_label = [''.join(map(str, x)) for x in val_loader.dataset.img_label]
        val_predict_label = predict(val_loader, model, 1)
        val_predict_label = np.vstack([
            val_predict_label[:, :11].argmax(1),
            val_predict_label[:, 11:22].argmax(1),
            val_predict_label[:, 22:33].argmax(1),
            val_predict_label[:, 33:44].argmax(1),
            val_predict_label[:, 44:55].argmax(1),
        ]).T
        val_label_pred = []
        for x in val_predict_label:
            val_label_pred.append(''.join(map(str, x[x != 10])))

        val_char_acc = np.mean(np.array(val_label_pred) == np.array(val_label))

        print('Epoch: {0}, Train loss: {1} \t Val loss: {2}'.format(epoch, train_loss, val_loss))
        print('Val Acc', val_char_acc)
    # 记录下验证集精度
        if val_loss < best_loss:
            best_loss = val_loss
        # print('Find better model in Epoch {0}, saving model.'.format(epoch))
            torch.save(model.state_dict(), './model.pt')

八、结果
最后正确率大概为55%
九、后续优化和目标
最终希望达到的正确率在90%，具体的后续优化方法暂时还没定。

十、最后的正确率
通过上传后，得到正确率只有0.3，这里的原因是，我进行了15个周期的训练，但是基础的baseline在10次周期左右就达到了不错的拟合效果，15个周期的训练应该已经过拟合了。