Datawhale-天池入门赛街景字符编码识别-Task5：模型集成

接上：Datawhale-天池入门赛街景字符编码识别-Task1：赛题理解、Datawhale-天池入门赛街景字符编码识别-Task2：数据读取与数据增强、Datawhale-天池入门赛街景字符编码识别-Task3：字符识别模型、Datawhale-天池入门赛街景字符编码识别-Task4：模型训练与验证

近期进展

将预测5个数字改为预测4个后，模型性能大幅提升，可见噪声对模型性能的影响。

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

    def __getitem__(self, idx):
        img = PIL.Image.open(self.img_path[idx]).convert('RGB')
        if self.transform:
            img = self.transform(img)
        # 将所有label处理为0-9+10（无数字）共11类的定长数字串
        lbl = np.array(self.img_label[idx], dtype=np.int)
        lbl = list(lbl) + (4 - len(lbl)) * [10]
        return img, torch.from_numpy(np.array(lbl[:4]))

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

模型也要相应改变。

class SVHN_Model2(nn.Module):
    def __init__(self):
        super(SVHN_Model2, self).__init__()
        model_conv = models.resnet50(pretrained=True)
        model_conv.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        model_conv = nn.Sequential(*list(model_conv.children())[:-1])
        self.cnn = model_conv
        # conv前用2d，fc前用1d，但是bn好像不能提升精度，只能提升训练速度，另外，bn和dropout不能一起用《Understanding the Disharmony Between Dropout and Batch Normalization by Variance Shift》
        # self.bn = nn.BatchNorm1d(2048)
        self.fc1 = nn.Linear(2048, 11)
        self.fc2 = nn.Linear(2048, 11)
        self.fc3 = nn.Linear(2048, 11)
        self.fc4 = nn.Linear(2048, 11)
    
    def forward(self, img):
        feat = self.cnn(img)
        feat = feat.view(feat.shape[0], -1)
        c1 = self.fc1(feat)
        c2 = self.fc2(feat)
        c3 = self.fc3(feat)
        c4 = self.fc4(feat)
        return c1, c2, c3, c4

各阶段也要相应变化。

def train(train_loader, model, criterion, optimizer):
    # 切换模型为训练模式
    model.train()
    train_loss = []
    
    for i, (input, target) in enumerate(train_loader):
        
        input = input.cuda()
        target = target.long().cuda()  

        c1, c2, c3, c4 = model(input)
        loss =  criterion(c1, target[:, 0]) + \
                criterion(c2, target[:, 1]) + \
                criterion(c3, target[:, 2]) + \
                criterion(c4, target[:, 3])
        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):
            input = input.cuda()
            target = target.long().cuda()    
            c1, c2, c3, c4 = model(input)
            loss =  criterion(c1, target[:, 0]) + \
                    criterion(c2, target[:, 1]) + \
                    criterion(c3, target[:, 2]) + \
                    criterion(c4, target[:, 3]) 
            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):
                input = input.cuda()
                c1, c2, c3, c4 = model(input)
                output = np.concatenate([
                    c1.data.cpu().numpy(), 
                    c2.data.cpu().numpy(),
                    c3.data.cpu().numpy(), 
                    c4.data.cpu().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

增加一个预测函数用于模型集成。

def predict2(test_loader, model1, model2, tta=10):
   
    model1.eval()
    model2.eval()
    test_pred_tta = True
    # TTA 次数
    for _ in range(tta):
        test_pred = []
        with torch.no_grad():
            for i, (input, target) in enumerate(test_loader):
                input = input.cuda()
                c1, c2, c3, c4= model1(input)
                output = np.concatenate([
                    c1.data.cpu().numpy(),
                    c2.data.cpu().numpy(),
                    c3.data.cpu().numpy(),
                    c4.data.cpu().numpy()], axis=1)
                c1, c2, c3, c4= model2(input)
                output2 = np.concatenate([
                    c1.data.cpu().numpy(),
                    c2.data.cpu().numpy(),
                    c3.data.cpu().numpy(),
                    c4.data.cpu().numpy()], axis=1)
                test_pred.append(output+output2)
                
            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

在训练阶段保存两个模型。

# 保存当前泛化性能最好的模型
        if val_loss < best_val:
            best_val = val_loss
            torch.save(model.state_dict(), time_path+'/best_val_model.pt')
            print("Save best val model at epoch:", epoch+1)

        # 保存当前最好模型
        if  val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), time_path+'/best_acc_model.pt')
            print("Save best acc model at epoch:", epoch+1)

最后集成两个模型的结果。

目前最好成绩

0.8591