third_party/SGMNet/train/train.py

import torch
import torch.optim as optim
from tqdm import trange
import os
from tensorboardX import SummaryWriter
import numpy as np
import cv2
from loss import SGMLoss, SGLoss
from valid import valid, dump_train_vis

import sys

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, ROOT_DIR)


from utils import train_utils


def train_step(optimizer, model, match_loss, data, step, pre_avg_loss):
    data["step"] = step
    result = model(data, test_mode=False)
    loss_res = match_loss.run(data, result)

    optimizer.zero_grad()
    loss_res["total_loss"].backward()
    # apply reduce on all record tensor
    for key in loss_res.keys():
        loss_res[key] = train_utils.reduce_tensor(loss_res[key], "mean")

    if loss_res["total_loss"] < 7 * pre_avg_loss or step < 200 or pre_avg_loss == 0:
        optimizer.step()
        unusual_loss = False
    else:
        optimizer.zero_grad()
        unusual_loss = True
    return loss_res, unusual_loss


def train(model, train_loader, valid_loader, config, model_config):
    model.train()
    optimizer = optim.Adam(model.parameters(), lr=config.train_lr)

    if config.model_name == "SGM":
        match_loss = SGMLoss(config, model_config)
    elif config.model_name == "SG":
        match_loss = SGLoss(config, model_config)
    else:
        raise NotImplementedError

    checkpoint_path = os.path.join(config.log_base, "checkpoint.pth")
    config.resume = os.path.isfile(checkpoint_path)
    if config.resume:
        if config.local_rank == 0:
            print("==> Resuming from checkpoint..")
        checkpoint = torch.load(
            checkpoint_path, map_location="cuda:{}".format(config.local_rank)
        )
        model.load_state_dict(checkpoint["state_dict"])
        best_acc = checkpoint["best_acc"]
        start_step = checkpoint["step"]
        optimizer.load_state_dict(checkpoint["optimizer"])
    else:
        best_acc = -1
        start_step = 0
    train_loader_iter = iter(train_loader)

    if config.local_rank == 0:
        writer = SummaryWriter(os.path.join(config.log_base, "log_file"))

    train_loader.sampler.set_epoch(
        start_step * config.train_batch_size // len(train_loader.dataset)
    )
    pre_avg_loss = 0

    progress_bar = (
        trange(start_step, config.train_iter, ncols=config.tqdm_width)
        if config.local_rank == 0
        else range(start_step, config.train_iter)
    )
    for step in progress_bar:
        try:
            train_data = next(train_loader_iter)
        except StopIteration:
            if config.local_rank == 0:
                print(
                    "epoch: ",
                    step * config.train_batch_size // len(train_loader.dataset),
                )
            train_loader.sampler.set_epoch(
                step * config.train_batch_size // len(train_loader.dataset)
            )
            train_loader_iter = iter(train_loader)
            train_data = next(train_loader_iter)

        train_data = train_utils.tocuda(train_data)
        lr = min(
            config.train_lr * config.decay_rate ** (step - config.decay_iter),
            config.train_lr,
        )
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr

        # run training
        loss_res, unusual_loss = train_step(
            optimizer, model, match_loss, train_data, step - start_step, pre_avg_loss
        )
        if (step - start_step) <= 200:
            pre_avg_loss = loss_res["total_loss"].data
        if (step - start_step) > 200 and not unusual_loss:
            pre_avg_loss = pre_avg_loss.data * 0.9 + loss_res["total_loss"].data * 0.1
        if unusual_loss and config.local_rank == 0:
            print(
                "unusual loss! pre_avg_loss: ",
                pre_avg_loss,
                "cur_loss: ",
                loss_res["total_loss"].data,
            )
        # log
        if config.local_rank == 0 and step % config.log_intv == 0 and not unusual_loss:
            writer.add_scalar("TotalLoss", loss_res["total_loss"], step)
            writer.add_scalar("CorrLoss", loss_res["loss_corr"], step)
            writer.add_scalar("InCorrLoss", loss_res["loss_incorr"], step)
            writer.add_scalar("dustbin", model.module.dustbin, step)

            if config.model_name == "SGM":
                writer.add_scalar("SeedConfLoss", loss_res["loss_seed_conf"], step)
                writer.add_scalar("MidCorrLoss", loss_res["loss_corr_mid"].sum(), step)
                writer.add_scalar(
                    "MidInCorrLoss", loss_res["loss_incorr_mid"].sum(), step
                )

        # valid ans save
        b_save = ((step + 1) % config.save_intv) == 0
        b_validate = ((step + 1) % config.val_intv) == 0
        if b_validate:
            (
                total_loss,
                acc_corr,
                acc_incorr,
                seed_precision_tower,
                seed_recall_tower,
                acc_mid,
            ) = valid(valid_loader, model, match_loss, config, model_config)
            if config.local_rank == 0:
                writer.add_scalar("ValidAcc", acc_corr, step)
                writer.add_scalar("ValidLoss", total_loss, step)

                if config.model_name == "SGM":
                    for i in range(len(seed_recall_tower)):
                        writer.add_scalar(
                            "seed_conf_pre_%d" % i, seed_precision_tower[i], step
                        )
                        writer.add_scalar(
                            "seed_conf_recall_%d" % i, seed_precision_tower[i], step
                        )
                    for i in range(len(acc_mid)):
                        writer.add_scalar("acc_mid%d" % i, acc_mid[i], step)
                    print(
                        "acc_corr: ",
                        acc_corr.data,
                        "acc_incorr: ",
                        acc_incorr.data,
                        "seed_conf_pre: ",
                        seed_precision_tower.mean().data,
                        "seed_conf_recall: ",
                        seed_recall_tower.mean().data,
                        "acc_mid: ",
                        acc_mid.mean().data,
                    )
                else:
                    print("acc_corr: ", acc_corr.data, "acc_incorr: ", acc_incorr.data)

                # saving best
                if acc_corr > best_acc:
                    print("Saving best model with va_res = {}".format(acc_corr))
                    best_acc = acc_corr
                    save_dict = {
                        "step": step + 1,
                        "state_dict": model.state_dict(),
                        "best_acc": best_acc,
                        "optimizer": optimizer.state_dict(),
                    }
                    save_dict.update(save_dict)
                    torch.save(
                        save_dict, os.path.join(config.log_base, "model_best.pth")
                    )

        if b_save:
            if config.local_rank == 0:
                save_dict = {
                    "step": step + 1,
                    "state_dict": model.state_dict(),
                    "best_acc": best_acc,
                    "optimizer": optimizer.state_dict(),
                }
                torch.save(save_dict, checkpoint_path)

            # draw match results
            model.eval()
            with torch.no_grad():
                if config.local_rank == 0:
                    if not os.path.exists(
                        os.path.join(config.train_vis_folder, "train_vis")
                    ):
                        os.mkdir(os.path.join(config.train_vis_folder, "train_vis"))
                    if not os.path.exists(
                        os.path.join(
                            config.train_vis_folder, "train_vis", config.log_base
                        )
                    ):
                        os.mkdir(
                            os.path.join(
                                config.train_vis_folder, "train_vis", config.log_base
                            )
                        )
                    os.mkdir(
                        os.path.join(
                            config.train_vis_folder,
                            "train_vis",
                            config.log_base,
                            str(step),
                        )
                    )
                res = model(train_data)
                dump_train_vis(res, train_data, step, config)
            model.train()

    if config.local_rank == 0:
        writer.close()