【阿里天池新人赛】之街景字符识别(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
下一章会对训练代码进行介绍。