语义分割实现地表建筑物识别5 模型训练与验证

语义分割实现地表建筑物识别5 模型训练与验证

学习目标

• 掌握数据划分方法和具体实践
• 掌握模型训练&模型调参过程

数据划分

• 训练集、验证集、测试集
• 划分方法：留出法、交叉验证法、自助采样法

模型调参流程

1. 初步构建简单的CNN模型，跑通训练、验证和预测的流程
2. 简单CNN模型的损失较大，尝试增加模型复杂度，并观察验证集精度
3. 在增加模型复杂度的同时增加数据扩增方法，直至验证集精度不变
   目前还处于baseline的跑通训练、验证阶段。
   调参部分需要根据训练误差、偏差判断需要增加模型复杂度还是通过数据扩增增加训练集数量、模型需要正则化。

使用Pytorch完成CNN的训练和验证

import torch.utils.data as D 

BATCH_SIZE = 10
EPOCHES = 5
best_loss = 10

# 划分训练集和验证集
dataset = TianChiDataset(
    train_mask['name'].values,
    train_mask['mask'].fillna('').values,
    trfm, False)

valid_idx, train_idx = [], []
for i in range(len(dataset)):
    if i % 7 == 0:
        valid_idx.append(i)
#     else:
    elif i % 7 == 1:
        train_idx.append(i)
        
train_ds = D.Subset(dataset, train_idx)
valid_ds = D.Subset(dataset, valid_idx)

# 构造训练集和验证集
train_loader = D.DataLoader(
    train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)

val_loader = D.DataLoader(
    valid_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
       
def train(train_loader,model,criterion,optimizer):
    # 切换模型为训练模式
    model.train()
    losses = []
    for image,target in tqdm_notebook(train_loader):
        # 正向传播
        image, target = image.to(DEVICE),target.float().to(DEVICE)
        optimizer.zero_grad()
        output = model(image)['out']        
        # 计算损失
        loss = criterion(output, target)
        # 反向传播
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    return losses

def validate(val_loader,model,criterion):
    # 切换模型为预测模型
    model.eval()
    val_loss = []

    # 不记录模型梯度信息
    with torch.no_grad():
        for i,(image,target) in enumerate(val_loader):
            # 正向传播
            image, target = image.to(DEVICE), target.float().to(DEVICE)
            output = model(image)['out']
            # 计算损失  
            loss = criterion(output, target) 
            losses.append(loss.item())
    return val_loss

# 模型训练与验证过程
model = Model1()
raw_line = '{:6d}' + '\u2502{:7.3f}'*2 + '\u2502{:6.2f}'
criterion = nn.CrossEntropyLoss(size_average=False)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-3)
best_loss = 1000.0
for epoch in range(1, EPOCHES+1):
    print('Epoch:',epoch)
    start_time = time.time()
    losses = train(train_loader,model,criterion,optimizer)
    val_loss = validate(val_loader,model,criterion)
    print(raw_line.format(epoch,np.array(losses).mean(),val_loss,
    (time.time()-start_time)/60**1))
    if val_loss < best_loss:
        best_loss = val_loss
        torch.save(model.state_dict(),'./model_best.pth')

## 保存最优验证集模型
torch.save(model_object.state_dict(),'model_best.pth')
model.load_state_dict(torch.load('model_best.pth'))

# 模型预测测试集结果
test_mask = pd.read_csv('./test_a_samplesubmit.csv', sep='\t', names=['name', 'mask'])
test_mask['name'] = test_mask['name'].apply(lambda x: './test_a/' + x)

for idx, name in enumerate(tqdm_notebook(test_mask['name'].iloc[:])):
    image = cv2.imread(name)
    image = trfm(image)
    with torch.no_grad():
        image = image.to(DEVICE)[None]
        score = model(image)['out'][0][0]
        score_sigmoid = score.sigmoid().cpu().numpy()
        score_sigmoid = (score_sigmoid > 0.5).astype(np.uint8)
        score_sigmoid = cv2.resize(score_sigmoid, (512, 512))
        
    subm.append([name.split('/')[-1], rle_encode(score_sigmoid)])