Train

$$Train$$

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.utils.data import DataLoader
from datetime import datetime
from major_dataset import LoadDataset
from major_evalution import eval_semantic_segmentation
import major_config

# ****************************************step1 数据处理**********************************************#
Load_train = LoadDataset([major_config.train_image, major_config.train_label], major_config.crop_size)
Load_val = LoadDataset([major_config.val_image, major_config.val_label], major_config.crop_size)

train_data = DataLoader(Load_train, batch_size=major_config.batchsize, shuffle=True, num_workers=1)
val_data = DataLoader(Load_val, batch_size=major_config.batchsize, shuffle=True, num_workers=1)

# *****************************************step2 模型*********************************************#
net = major_config.model
net = net.to(major_config.device)

# ******************************************step3 损失函数********************************************#
criterion = nn.NLLLoss().to(major_config.device)  # NLLLoss有利于最后激活层的替换

# ******************************************step4 优化器********************************************#
optimizer = optim.Adam(net.parameters(), lr=1e-4)

# ******************************************step5 训练********************************************#
def train(model):
    best = [0]  # 存储最优指标，用于Early Stopping
    net = model.train()  # 指定模型为训练模式，即可以进行参数更新
    # 训练轮次
    for epoch in range(major_config.num_epoch):
        print('Epoch is [{}/{}]'.format(epoch + 1, major_config.num_epoch))
        # 每20次epoch,lr学习率降一半
        if epoch % 20 == 0 and epoch != 0:
            for group in optimizer.param_groups:
                group['lr'] *= 0.5
        # 指标初始化
        train_loss = 0
        train_pa = 0
        train_mpa = 0
        train_miou = 0
        train_fwiou = 0
        # 训练批次
        for i, sample in enumerate(train_data):
            # 载入数据
            img_data = sample['img'].to(major_config.device)
            img_label = sample['label'].to(major_config.device)
            # 训练
            out = net(img_data)
            out = F.log_softmax(out, dim=1)
            loss = criterion(out, img_label)  # loss计算
            optimizer.zero_grad()  # 需要梯度清零，再反向传播
            loss.backward()  # 反向传播
            optimizer.step()  # 参数更新
            train_loss += loss.item()  # loss累加
            # 评估
            # 预测值
            pre_label = out.max(dim=1)[1].data.cpu().numpy()  # [1]：表示返回索引
            pre_label = [i for i in pre_label]
            # 真实值
            true_label = img_label.data.cpu().numpy()
            true_label = [i for i in true_label]
            # 计算所有的评价指标
            eval_metrix = eval_semantic_segmentation(pre_label, true_label)
            # 各评价指标计算
            train_pa += eval_metrix['pa']
            train_mpa += eval_metrix['mpa']
            train_miou += eval_metrix['miou']
            train_fwiou += eval_metrix['fwiou']
            #  打印损失
            print('|batch[{}/{}]|batch_loss {: .8f}|'.format(i + 1, len(train_data), loss.item()))
        #  评价指标打印格式定义
        metric_description = '|Train PA|: {:.5f}|\n|Train MPA|: {:.5f}|\n|Train MIou|: {:.5f}|\n|Train FWIou|: {:.5f}|'.format(
            train_pa / len(train_data),
            train_mpa / len(train_data),
            train_miou / len(train_data),
            train_fwiou / len(train_data),
        )
        #  打印评价指标
        print(metric_description)
        #  根据train_miou，保存最优模型
        if max(best) <= train_miou / len(train_data):
            best.append(train_miou / len(train_data))
            torch.save(net.state_dict(), major_config.path_saved_model)

# ******************************************step6 评价********************************************#
def evaluate(model):
    net = model.eval()
    eval_loss = 0
    eval_acc = 0
    eval_miou = 0
    eval_class_acc = 0

    prec_time = datetime.now()
    for j, sample in enumerate(val_data):
        valImg = sample['img'].to(major_config.device)
        valLabel = sample['label'].long().to(major_config.device)

        out = net(valImg)
        out = F.log_softmax(out, dim=1)
        loss = criterion(out, valLabel)
        eval_loss = loss.item() + eval_loss
        pre_label = out.max(dim=1)[1].data.cpu().numpy()
        pre_label = [i for i in pre_label]

        true_label = valLabel.data.cpu().numpy()
        true_label = [i for i in true_label]

        eval_metrics = eval_semantic_segmentation(pre_label, true_label)
        eval_acc = eval_metrics['mean_class_accuracy'] + eval_acc
        eval_miou = eval_metrics['miou'] + eval_miou

    cur_time = datetime.now()
    h, remainder = divmod((cur_time - prec_time).seconds, 3600)
    m, s = divmod(remainder, 60)
    time_str = 'Time: {:.0f}:{:.0f}:{:.0f}'.format(h, m, s)

    val_str = ('|Valid Loss|: {:.5f} \n|Valid Acc|: {:.5f} \n|Valid Mean IU|: {:.5f} \n|Valid Class Acc|:{:}'.format(
        eval_loss / len(train_data),
        eval_acc / len(val_data),
        eval_miou / len(val_data),
        eval_class_acc / len(val_data)))
    print(val_str)
    print(time_str)


if __name__ == "__main__":
    train(net)
    # evaluate(net) 验证可以自己设置每训练多少次，验证一次，所以，evaluate()函数可以放到train()里面