上篇文章复现了vgg16模型，本文用vgg16来训练一个用于图像分类任务的baseline，因为本文是针对华为2022全球校园AI算法精英大赛的，所以先介绍下比赛。

比赛简介

比赛的主办方是华为，一共有三个赛道：

• 知识驱动口语对话
• 车道渲染数据智能质检
• 广告-信息流跨域ctr预估

看上面的介绍，可以分别对应NLP、CV、推荐算法三种类型。感兴趣的同学可以自由参加，本次比赛必须是在校学生才可以参与，相对来说可能竞争要小一点。

赛题2简介

报名了赛题2车道渲染数据智能质检，赛题简介：

本赛题为选手提供地图渲染数据及其部分标注。训练集包括车道渲染数据的图片集及部分标注，选手使用训练数据进行模型训练。测试集分为A/B两个测试集。测试集仅提供车道渲染数据图片集，选手使用模型预测测试图片是否存在问题。

从上面可以发现这其实是一个对图像进行二分类的问题，预测的结果就是0（没有问题）和1（有问题）。

baseline

赛题的要解决的问题刚好可以用上篇文章推送的vgg16来解决，本文提供一个baseline，只不过使用pytorch中提供的预训练好的模型，节省训练时间。

数据分析

工欲善其事必先利其器。对于深度学习来说，数据就是一切的基础，所以一定要先分析训练数据。 赛事方提供的数据包括：

• train_image.rar：训练集图片，png格式
• train_label.rar：训练集标注，csv格式
• test_image.rar：测试集图片，png格式

也就是说我们要利用已经提供的图片和csv文件训练一个模式，然后这个模型接受一张图片，可以输出这张图片"是否有问题"这样子一个结果。

因为原始的数据集太大了（三十多个g），所以本文训练数据集使用Datawhale提供的从原始数据抽样出来的一小部分数据集（1g多），这里表示感谢，大家可以加入Datawhale关于本赛题的交流群，具体见参考链接1

训练数据的图片大概是这个样子：

就是车道渲染数据，看起来感觉是用在导航的过程中提示司机怎么走。。。

代码

下面直接上代码。还是按照之前的分为数据、模型、损失函数、优化器和迭代训练5个部分。

导入包

# 数据
import torch
import cv2
import PIL
import pandas as pd
import numpy as np
import glob
import torchvision.transforms as transforms
# 模型
import torch.nn as nn
import torchvision.models as models
import torch.utils.data as Data
# 损失函数和优化
import torch.optim as optim
# 训练
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
# 其他
import gc
import os

设置一下gou

device = 'cuda' if torch.cuda.is_available() else 'cpu'
if device == 'cuda':
    gc.collect()
    torch.cuda.empty_cache()

训练数据制作

# 在label文件中，有些存在问题的label可能是多个，所以加上names参数
# 但是因为是二分类问题，只要知道其不是1就行，不在乎其存在几个问题，所以只取前两列
train_data_df = pd.read_csv('./train_label.csv',names=list(range(5)))
train_data_df = train_data_df.iloc[:,:2]

train_data_df['image'] = train_data_df[0].apply(lambda x:'./train_image/'+x)
train_data_df['label'] = train_data_df[1].apply(lambda x: 1 if x!= 0 else 0)
train_data_df  = train_data_df[['image','label']]

class RoadData(Data.Dataset):
    def __init__(self, df, transform=None):
        super(RoadData, self).__init__
        self.df = df
        self.transform = transform

    def __getitem__(self, idx):
        image = self.df.loc[idx,'image']
        try:
            image = PIL.Image.open(image).convert('RGB')
        except:
            # 因为有些图像坏了，读不出来
            image = np.zeros((256,256), dtype=np.uint8)
            image = PIL.Image.fromarray(image).convert('RGB')

        label = self.df.loc[idx,'label']
        label = torch.tensor(label)

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

        return image, label

    def __len__(self):

        return len(self.df)
transform = transforms.Compose([
    transforms.Resize([224,224]),
    transforms.ToTensor(), 
])

train_ds = RoadData(df=train_data_df, transform= transform)
train_dl = Data.DataLoader(train_ds, batch_size=16, shuffle=True, num_workers=0)

模型

model = models.vgg16(pretrained = True) # 加载预训练的vgg16模型
model.classifier[6] = nn.Linear(4096,2)
model.to(device)

损失函数和优化器

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum =0.9)

迭代训练

def train_one_epoch(dataloader, model, criterion, optimizer):
    model.train()

    running_loss = 0.0
    running_accuracy = 0.0
    total = 0

    for idx, data in enumerate(dataloader):
        images, labels = data
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        outs = model(images)
        loss = criterion(outs,labels)
        loss.backward()
        optimizer.step()

        running_loss += loss
        _, preds = torch.max(outs,1)
        running_accuracy += (preds==labels).sum().item()
        total += images.shape[0]

    train_loss = running_loss / total
    train_metric = running_accuracy /total

    return train_loss, train_metric        

epochs = 10
losses = []
metrics = []

for epoch in range(epochs):
    train_loss, train_metric = train_one_epoch(model=model, dataloader=train_dl, criterion=criterion, optimizer=optimizer)
    losses.append(train_loss)
    metrics.append(train_metric)
    print(f"epoch:{epoch}, loss:{train_loss}, metrics:{train_metric}")

预测

def predict_one_image(data, model, transform=transform):
    model.eval()

    try:
        image = PIL.Image.open(data).convert('RGB')
    except:
        # 因为有些图像坏了，读不出来
        image = np.zeros((256,256), dtype=np.uint8)
        image = PIL.Image.fromarray(image).convert('RGB')
    image = transform(image).unsqueeze(0).to(device)
    out = model(image)
    prob = nn.Softmax(dim=1)
    out = prob(out)
    out = out.data.cpu().numpy()[0,1] # 输出有问题的概率
    
    return out

test_df = pd.DataFrame({"image":glob.glob(r'./test_images/*.png')})
test_df['imagename'] = test_df['image'].apply(lambda x: os.path.basename(x))

test_df['defect_prob'] = test_df['image'].apply(predict_one_image,model=model,transform=transform)
submission = test_df[['imagename','defect_prob']]
submission.to_csv('./submission.csv', index = False)

然后就得到对应的submission文件啦，就可以提交了！

寻求帮助

按照上面的过程生成了文件，但是我在提交的时候报错了。。。

我确认流程没有错，test文件夹里面也只有4000张图片，带上列名，一共4001行。。。明明没有错，却让我提交10001行，难道华为跟电信有什么猫腻？

求大佬帮忙解决。。

虽然自己没提交成功，baseline放出来，希望大家能用上~

更新

原来是官方更新了数据，最新的数据中test数据集一共是1w张。。。但是不知道为什么我下载的还是4k张。。。无语ing

参考

【1】https://mp.weixin.qq.com/s/PXNleJaLeNpnGvulxG3WAQ