【阿里天池新人赛】之街景字符识别（2）

【阿里天池新人赛】之街景字符识别（2）

正式赛时间：2020.5.14-2020.6.24
比赛网址：https://tianchi.aliyun.com/competition/entrance/531795/introduction

1.数据读取（dataset.py）

由于本次赛题我们使用Pytorch框架，接下来将是解决赛题的第一步使用Pytorch读取赛题数据。
在Pytorch中数据是通过Dataset进行封装，并通过DataLoder进行并行读取。所以我们只需要重载一下数据读取的逻辑就可以完成数据的读取。

from torch.utils.data.dataset import Dataset
import torch
from PIL import Image
import numpy as np


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)

        # 设置最⻓的字符⻓度为5个
        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.构建模型（model.py）

这里我们采用预先训练ResNet 18网络构建。

import torch.nn as nn
import torchvision.models as models


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

        model_conv = models.resnet18(pretrained=True)
        # print(model_conv)
        model_conv.avgpool = nn.AdaptiveAvgPool2d(1)
        num_ftrs = model_conv.fc.in_features
        model_conv = nn.Sequential(*list(model_conv.children())[:-1])
        self.cnn = model_conv
        self.fc1 = nn.Linear(num_ftrs, 11)
        self.fc2 = nn.Linear(num_ftrs, 11)
        self.fc3 = nn.Linear(num_ftrs, 11)
        self.fc4 = nn.Linear(num_ftrs, 11)
        self.fc5 = nn.Linear(num_ftrs, 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

下一章会对训练代码进行介绍。