比赛链接：https://tianchi.aliyun.com/competition/entrance/531795/introduction

baseline思路：使用CNN进行定长字符分类；

运行系统要求：Python2/3，内存4G，有无GPU都可以

线上得分取决于训练轮数，假设比赛数据路径为../input

import os, sys, glob, shutil, json
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import cv2from PIL import Image
import numpy as npfrom tqdm import tqdm, tqdm_notebook%pylab inlineimport torch
torch.manual_seed(0)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = Trueimport torchvision.models as models
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data.dataset import Dataset

Populating the interactive namespace from numpy and matplotlib

定义读取数据集

class SVHNDataset(Dataset):def __init__(self, img_path, img_label, transform=None):self.img_path = img_pathself.img_label = img_label if transform is not None:self.transform = transformelse:self.transform = Nonedef __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)

定义读取数据dataloader

假设数据存放在../input文件夹下，并进行解压。

train_path = glob.glob('../input/train/*.png')
train_path.sort()
train_json = json.load(open('../input/train.json'))
train_label = [train_json[x]['label'] for x in train_json]
print(len(train_path), len(train_label))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=40, shuffle=True, num_workers=10,
)val_path = glob.glob('../input/val/*.png')
val_path.sort()
val_json = json.load(open('../input/val.json'))
val_label = [val_json[x]['label'] for x in val_json]
print(len(val_path), len(val_label))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=40, shuffle=False, num_workers=10,
)

30000 3000010000 10000

定义分类模型

这里使用ResNet18的模型进行特征提取

class SVHN_Model1(nn.Module):def __init__(self):super(SVHN_Model1, self).__init__()model_conv = models.resnet18(pretrained=True)model_conv.avgpool = nn.AdaptiveAvgPool2d(1)model_conv = nn.Sequential(*list(model_conv.children())[:-1])self.cnn = model_convself.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]) + \criterion(c1, target[:, 1]) + \criterion(c2, target[:, 2]) + \criterion(c3, target[:, 3]) + \criterion(c4, target[:, 4])# loss /= 6optimizer.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]) + \criterion(c1, target[:, 1]) + \criterion(c2, target[:, 2]) + \criterion(c3, target[:, 3]) + \criterion(c4, target[:, 4])# loss /= 6val_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_predelse:test_pred_tta += test_predreturn 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
if use_cuda:model = model.cuda()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),]).Tval_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')

预测并生成提交文件

test_path = glob.glob('../input/test_a/*.png')
test_path.sort()
test_json = json.load(open('../input/test_a.json'))
test_label = [[1]] * len(test_path)
print(len(test_path), len(test_label))test_loader = torch.utils.data.DataLoader(SVHNDataset(test_path, test_label,transforms.Compose([transforms.Resize((70, 140)),# transforms.RandomCrop((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=40, shuffle=False, num_workers=10,
)

40000 40000

# 加载保存的最优模型
model.load_state_dict(torch.load('model.pt'))test_predict_label = predict(test_loader, model, 1)
print(test_predict_label.shape)test_label = [''.join(map(str, x)) for x in test_loader.dataset.img_label]
test_predict_label = np.vstack([test_predict_label[:, :11].argmax(1),test_predict_label[:, 11:22].argmax(1),test_predict_label[:, 22:33].argmax(1),test_predict_label[:, 33:44].argmax(1),test_predict_label[:, 44:55].argmax(1),
]).Ttest_label_pred = []
for x in test_predict_label:test_label_pred.append(''.join(map(str, x[x!=10])))import pandas as pd
df_submit = pd.read_csv('../input/test_A_sample_submit.csv')
df_submit['file_code'] = test_label_pred
df_submit.to_csv('submit.csv', index=None)

Datawhale 零基础入门CV赛事-Baseline

— By:阿水

天池数据大神，CV爱好者。 公众号：Coggle数据科学 知乎：https://www.zhihu.com/people/finlayliu

— By:小武

南昌航空大学研究生，CV爱好者。 博客：https://blog.csdn.net/weixin_40647819

因为第一次在天池上打比赛，也是第一次参加比赛，数据集都不知道怎么下载，搞了好长时间才按照 baseline 一步步下来，网速慢下载数据真是艰辛，本周要毕业答辩了，也没那么多时间来研究，草草的完成任务吧。。。。