code/libs/utils/train_utils.py

import os
import shutil
import time
import pickle
import json
import copy
import collections
from terminaltables import AsciiTable

import numpy as np
import random
from copy import deepcopy

import torch
import torch.optim as optim
import torch.backends.cudnn as cudnn

from .lr_schedulers import LinearWarmupMultiStepLR, LinearWarmupCosineAnnealingLR
from .postprocessing import postprocess_results
from .detect_eval import compute_AP_AR
from ..modeling import MaskedConv1D, Scale, AffineDropPath, LayerNorm
from IPython import embed
from tqdm import tqdm

################################################################################
def calculate_IoU_batch2(i0, i1):
    union = (np.min(np.stack([i0[0], i1[0]], 0), 0), np.max(np.stack([i0[1], i1[1]], 0), 0))
    inter = (np.max(np.stack([i0[0], i1[0]], 0), 0), np.min(np.stack([i0[1], i1[1]], 0), 0))
    iou = 1.0 * (inter[1] - inter[0] + 1) / (union[1] - union[0] + 1)
    iou[union[1] - union[0] < -1e-5] = 0
    iou[iou < 0] = 0.0
    return iou

j_list = [70, 80, 90]
IOU_THRESHOLDS = [0.75, 0.85, 0.95]
AR_IOU_THRESHOLDS = [0.75, 0.85, 0.95]
print(f'AR_IOU_THRESHOLDS: {AR_IOU_THRESHOLDS}', flush=True)
MAX_DETECTIONS = [1, 5, 10]

# [nb, 2], [nb, 2]
def top_n_metric(preds, label):
    result = {}
    bsz = preds[0].shape[0]
    top_iou = []
    for pred in preds:
        iou = calculate_IoU_batch2((pred[:, 0], pred[:, 1]), (label[:, 0], label[:, 1]))
        top_iou.append(iou)
    iou = np.max(np.stack(top_iou, 1), 1)
    result['mIoU'] = np.mean(iou)
    # for i in range(1, 10, 1):
    #     result['IoU@0.{}'.format(i)] = 1.0 * np.sum(iou >= i / 10) / bsz
    for i in np.arange(0, 100, 1):
        result['IoU@{}'.format(i / 100)] = np.sum(iou >= i / 100) / bsz
    return result

def nms_temporal(x1,x2,s, overlap):
    pick = []
    assert len(x1)==len(s)
    assert len(x2)==len(s)
    if len(x1)==0:
        return pick
    #     union = map(operator.sub, x2, x1) # union = x2-x1
    union = list(map(operator.sub, x2, x1)) # union = x2-x1

    I = [i[0] for i in sorted(enumerate(s), key=lambda x:x[1])] # sort and get index

    while len(I)>0:
        i = I[-1]
        pick.append(i)
        xx1 = [max(x1[i],x1[j]) for j in I[:-1]]
        xx2 = [min(x2[i],x2[j]) for j in I[:-1]]
        inter = [max(0.0, k2-k1) for k1, k2 in zip(xx1, xx2)]
        o = [inter[u]/(union[i] + union[I[u]] - inter[u]) for u in range(len(I)-1)]
        I_new = []
        for j in range(len(o)):
            if o[j] <=overlap:
                I_new.append(I[j])
        I = I_new
    return pick

#     table = AsciiTable(display_list)
#     return table.table

def display_python_performance(x, data='CharadesSTA'):
    """
    """
    global IOU_THRESHOLDS, MAX_DETECTIONS

    ap_keys = [f"mAP@{t}".rstrip("0").rstrip(".") for t in IOU_THRESHOLDS] + ["mAP"]
    ar_keys = [f"AR@{k}" for k in MAX_DETECTIONS] + ["mAR"]

    # ===== AP 部分 =====
    ap_table_data = [["Type"] + ap_keys]
    ap_row = ["AP"]
    for k in ap_keys:
        v = getattr(x[k], "avg", x[k]) if hasattr(x[k], "avg") else x[k]
        ap_row.append(f"{v * 100:.2f}")
    ap_table_data.append(ap_row)

    ap_table = AsciiTable(ap_table_data)
    ap_str = ap_table.table

    # ===== AR 部分 =====
    ar_table_data = [["Type"] + ar_keys]
    ar_row = ["AR"]
    for k in ar_keys:
        v = getattr(x[k], "avg", x[k]) if hasattr(x[k], "avg") else x[k]
        ar_row.append(f"{v * 100:.2f}")
    ar_table_data.append(ar_row)

    ar_table = AsciiTable(ar_table_data)
    ar_str = ar_table.table

    final_output = f"{ap_str}\n{ar_str}"

    return final_output



def print_md_performance(x, data='CharadesSTA'):
    item = ''
    i_list = [1]
    # j_list = [7, 8, 9]

    for i in i_list:
        for j in j_list:
            key = 'IoU@{}'.format(j/100)
            item += '| {:.2f} '.format(x[key].avg * 100)
        item += '| {:.2f} |\n'.format(x[f'mIoU'].avg * 100)

    for i in i_list:
        for j in j_list:
            key = 'IoU@{}'.format(j/100)
            item += '&{:.2f} '.format(x[key].avg * 100)
        item += '&{:.2f} |'.format(x[f'mIoU'].avg * 100)
        if i_list[-1] != i:
            item += '\n'
    return item

def get_average_performance(x, data='CharadesSTA', i=1):
    # # j_list = [7, 8, 9]
    # average = 0.0
    # for j in j_list:
    #     key = 'IoU@{}'.format(j / 100)
    #     average += x[key].avg
    # average = average / len(j_list)
    #
    # # average = x['IoU@0.9'].avg
    # average *= 100.0
    # return average
    if isinstance(x['mAP'], AverageMeter):
        average = x['mAP'].avg
    else:
        average = x['mAP']
    average *= 100.0
    return average

def print_python_performance(x, data='CharadesSTA'):
    item = ''
    i_list = [1]
    # j_list = [7, 8, 9]

    item += '{}:  {:.2f}\n'.format('average_R1', get_average_performance(x, data))
    item += '{}:        {:.2f}\n'.format('mIoU', x['mIoU'].avg * 100)
    for i in i_list:
        for j in j_list:
            key = 'IoU@{}'.format(j/100)
            item += '{}: {:.2f}  '.format(key, x[key].avg * 100)
        # if i == 1:
        #     item += '\n'
    return  item

def merge_ResultSaveObj(ResultSaveObj_list):
    ret = ResultSaveObj()
    ret.merge(ResultSaveObj_list)
    return ret

class ResultSaveObj(object):
    def __init__(self):
        self.save_flag = False
        self.video_info = collections.defaultdict(list)

    def merge(self, ResultSaveObj_list):
        for itm in ResultSaveObj_list:
            for k, v in itm.video_info.items():
                # self.video_info[k] += copy.deepcopy(v)
                self.video_info[k] += v
            # self.video_info += copy.deepcopy(itm.video_info)

    def to_numpy(self, x):
        return x.detach().cpu().numpy()

    def add_batch(self, save_result):
        # if not self.save_flag:
        #     return
        for k, v in save_result.items():
            self.video_info[k].append(v)

    def eval(self):
        gt_time = self.video_info['gt_time']
        pred_time = self.video_info['pred_time']
        score = self.video_info['score']
        # embed()
        metric_dict_1 = compute_AP_AR(
            pred_time, gt_time, score,
            iou_thresholds_ap=np.array(IOU_THRESHOLDS), iou_thresholds_ar=np.array(AR_IOU_THRESHOLDS), ar_points=MAX_DETECTIONS
        )
        return metric_dict_1


def fix_random_seed(seed, include_cuda=True):
    rng_generator = torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    if include_cuda:
        # training: disable cudnn benchmark to ensure the reproducibility
        cudnn.enabled = True
        cudnn.benchmark = False
        cudnn.deterministic = True
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        # this is needed for CUDA >= 10.2
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
        torch.use_deterministic_algorithms(True)
    else:
        cudnn.enabled = True
        cudnn.benchmark = True
    return rng_generator


# def save_checkpoint(state, is_best, file_folder,
def save_checkpoint(state, file_folder,
                    file_name='checkpoint.pth.tar'):
    pass

def print_model_params(model):
    for name, param in model.named_parameters():
        print(name, param.min().item(), param.max().item(), param.mean().item())
    return


def make_optimizer(model, optimizer_config):
    """create optimizer
    return a supported optimizer
    """
    # separate out all parameters that with / without weight decay
    # see https://github.com/karpathy/minGPT/blob/master/mingpt/model.py#L134
    decay = set()
    no_decay = set()
    whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv1d, MaskedConv1D)
    blacklist_weight_modules = (LayerNorm, torch.nn.GroupNorm)

    # loop over all modules / params
    for mn, m in model.named_modules():
        for pn, p in m.named_parameters():
            fpn = '%s.%s' % (mn, pn) if mn else pn # full param name
            if pn.endswith('bias'):
                # all biases will not be decayed
                no_decay.add(fpn)
            elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
                # weights of whitelist modules will be weight decayed
                decay.add(fpn)
            elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
                # weights of blacklist modules will NOT be weight decayed
                no_decay.add(fpn)
            elif pn.endswith('scale') and isinstance(m, (Scale, AffineDropPath)):
                # corner case of our scale layer
                no_decay.add(fpn)
            elif pn.endswith('rel_pe'):
                # corner case for relative position encoding
                no_decay.add(fpn)

    # validate that we considered every parameter
    param_dict = {pn: p for pn, p in model.named_parameters()}
    inter_params = decay & no_decay
    union_params = decay | no_decay
    assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), )
    assert len(param_dict.keys() - union_params) == 0, \
        "parameters %s were not separated into either decay/no_decay set!" \
        % (str(param_dict.keys() - union_params), )

    # create the pytorch optimizer object
    optim_groups = [
        {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": optimizer_config['weight_decay']},
        {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
    ]

    if optimizer_config["type"] == "SGD":
        optimizer = optim.SGD(
            optim_groups,
            lr=optimizer_config["learning_rate"],
            momentum=optimizer_config["momentum"]
        )
    elif optimizer_config["type"] == "AdamW":
        optimizer = optim.AdamW(
            optim_groups,
            lr=optimizer_config["learning_rate"]
        )
    else:
        raise TypeError("Unsupported optimizer!")

    return optimizer


def make_scheduler(
    optimizer,
    optimizer_config,
    num_iters_per_epoch,
    last_epoch=-1
):
    """create scheduler
    return a supported scheduler
    All scheduler returned by this function should step every iteration
    """
    if optimizer_config["warmup"]:
        max_epochs = optimizer_config["epochs"] + optimizer_config["warmup_epochs"]
        max_steps = max_epochs * num_iters_per_epoch

        # get warmup params
        warmup_epochs = optimizer_config["warmup_epochs"]
        warmup_steps = warmup_epochs * num_iters_per_epoch

        # with linear warmup: call our custom schedulers
        if optimizer_config["schedule_type"] == "cosine":
            # Cosine
            scheduler = LinearWarmupCosineAnnealingLR(
                optimizer,
                warmup_steps,
                max_steps,
                last_epoch=last_epoch
            )

        elif optimizer_config["schedule_type"] == "multistep":
            # Multi step
            steps = [num_iters_per_epoch * step for step in optimizer_config["schedule_steps"]]
            scheduler = LinearWarmupMultiStepLR(
                optimizer,
                warmup_steps,
                steps,
                gamma=optimizer_config["schedule_gamma"],
                last_epoch=last_epoch
            )
        else:
            raise TypeError("Unsupported scheduler!")

    else:
        max_epochs = optimizer_config["epochs"]
        max_steps = max_epochs * num_iters_per_epoch

        # without warmup: call default schedulers
        if optimizer_config["schedule_type"] == "cosine":
            # step per iteration
            scheduler = optim.lr_scheduler.CosineAnnealingLR(
                optimizer,
                max_steps,
                last_epoch=last_epoch
            )

        elif optimizer_config["schedule_type"] == "multistep":
            # step every some epochs
            steps = [num_iters_per_epoch * step for step in optimizer_config["schedule_steps"]]
            scheduler = optim.lr_scheduler.MultiStepLR(
                optimizer,
                steps,
                gamma=schedule_config["gamma"],
                last_epoch=last_epoch
            )
        else:
            raise TypeError("Unsupported scheduler!")

    return scheduler


class AverageMeter(object):
    """Computes and stores the average and current value.
    Used to compute dataset stats from mini-batches
    """
    def __init__(self):
        self.initialized = False
        self.val = None
        self.avg = None
        self.sum = None
        self.count = 0.0

    def initialize(self, val, n):
        self.val = val
        self.avg = val
        self.sum = val * n
        self.count = n
        self.initialized = True

    def update(self, val, n=1):
        if not self.initialized:
            self.initialize(val, n)
        else:
            self.add(val, n)

    def add(self, val, n):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


class ModelEma(torch.nn.Module):
    def __init__(self, model, decay=0.999, device=None):
        super().__init__()
        # make a copy of the model for accumulating moving average of weights
        self.module = deepcopy(model)
        self.module.eval()
        self.decay = decay
        self.device = device  # perform ema on different device from model if set
        if self.device is not None:
            self.module.to(device=device)

    def _update(self, model, update_fn):
        with torch.no_grad():
            for ema_v, model_v in zip(self.module.state_dict().values(), model.state_dict().values()):
                if self.device is not None:
                    model_v = model_v.to(device=self.device)
                ema_v.copy_(update_fn(ema_v, model_v))

    def update(self, model):
        self._update(model, update_fn=lambda e, m: self.decay * e + (1. - self.decay) * m)

    def set(self, model):
        self._update(model, update_fn=lambda e, m: m)


################################################################################
def train_one_epoch(
    train_loader,
    model,
    optimizer,
    scheduler,
    curr_epoch,
    model_ema = None,
    clip_grad_l2norm = -1,
    tb_writer = None,
    print_freq = 20
):
    """Training the model for one epoch"""
    # set up meters
    batch_time = AverageMeter()
    losses_tracker = {}
    # number of iterations per epoch
    num_iters = len(train_loader)
    # switch to train mode
    model.train()

    # main training loop
    print()
    print('='*100)
    print("[Train]: Epoch {:d} started".format(curr_epoch))
    print('=' * 100)
    start = time.time()
    for iter_idx, video_list in enumerate(train_loader, 0):
        # zero out optim
        optimizer.zero_grad(set_to_none=True)
        # forward / backward the model
        losses = model(video_list)
        losses['final_loss'].backward()
        # gradient cliping (to stabilize training if necessary)
        if clip_grad_l2norm > 0.0:
            torch.nn.utils.clip_grad_norm_(
                model.parameters(),
                clip_grad_l2norm
            )
        # step optimizer / scheduler
        optimizer.step()
        scheduler.step()

        if model_ema is not None:
            model_ema.update(model)

        # printing (only check the stats when necessary to avoid extra cost)
        if (iter_idx != 0) and (iter_idx % print_freq) == 0:
            # measure elapsed time (sync all kernels)
            torch.cuda.synchronize()
            batch_time.update((time.time() - start) / print_freq)
            start = time.time()

            # track all losses
            for key, value in losses.items():
                # init meter if necessary
                if key not in losses_tracker:
                    losses_tracker[key] = AverageMeter()
                # update
                losses_tracker[key].update(value.item())

            # log to tensor board
            lr = scheduler.get_last_lr()[0]
            global_step = curr_epoch * num_iters + iter_idx
            if tb_writer is not None:
                # learning rate (after stepping)
                tb_writer.add_scalar(
                    'train/learning_rate',
                    lr,
                    global_step
                )
                # all losses
                tag_dict = {}
                for key, value in losses_tracker.items():
                    if key != "final_loss":
                        tag_dict[key] = value.val
                tb_writer.add_scalars(
                    'train/all_losses',
                    tag_dict,
                    global_step
                )
                # final loss
                tb_writer.add_scalar(
                    'train/final_loss',
                    losses_tracker['final_loss'].val,
                    global_step
                )

            # print to terminal
            block1 = 'Epoch: [{:03d}][{:05d}/{:05d}]'.format(
                curr_epoch, iter_idx, num_iters
            )
            block2 = 'Time {:.1f} ({:.1f})'.format(
                batch_time.val, batch_time.avg
            )
            block3 = 'Loss {:.3f} ({:.3f})'.format(
                losses_tracker['final_loss'].val,
                losses_tracker['final_loss'].avg
            )
            block4 = ''
            for key, value in losses_tracker.items():
                if key != "final_loss":
                    block4  += '\t{:s} {:.3f} ({:.3f})'.format(
                        key, value.val, value.avg
                    )

            print(f'{block1} {block2}\n{block3} {block4}\n')

    # finish up and print
    lr = scheduler.get_last_lr()[0]
    print("\n[Train]: Epoch {:d} finished with lr={:.8f}\n\n".format(curr_epoch, lr))
    return


def valid_one_epoch(
    val_loader,
    model,
    curr_epoch,
    ext_score_file = None,
    evaluator = None,
    output_file = None,
    tb_writer = None,
    print_freq = 20,
):
    """Test the model on the validation set"""
    # either evaluate the results or save the results
    # assert (evaluator is not None) or (output_file is not None)

    # set up meters
    batch_time = AverageMeter()
    # switch to evaluate mode
    model.eval()
    # dict for results (for our evaluation code)
    results = {
        'video-id': [],
        't-start' : [],
        't-end': [],
        'label': [],
        'score': []
    }

    # loop over validation set
    start = time.time()
    acc_metrics_logger = collections.defaultdict(lambda: AverageMeter())
    result_save_obj_dict = {}

    result_save_obj = ResultSaveObj()
    iter_list = enumerate(val_loader, 0)
    batch_time = AverageMeter()

    for iter_idx, video_list in iter_list:
        # forward the model (wo. grad)
        with torch.no_grad():
            output_with_pred_list = model(video_list)

            # unpack the results into ANet format
            local_weight_list = [0.0]
            output_with_pred_list = [output_with_pred_list]
            local_weight_id = -1
            for output, local_weight in zip(output_with_pred_list, local_weight_list):
                local_weight_id += 1
                if local_weight not in result_save_obj_dict:
                    result_save_obj = ResultSaveObj()
                    result_save_obj_dict[local_weight] = result_save_obj
                else:
                    result_save_obj = result_save_obj_dict[local_weight]
                num_vids = len(output)
                for vid_idx in range(num_vids):
                    if output[vid_idx]['segments'].shape[0] > 0:
                        results['video-id'].extend(
                            [output[vid_idx]['video_id']] *
                            output[vid_idx]['segments'].shape[0]
                        )
                        results['t-start'].append(output[vid_idx]['segments'][:, 0])
                        results['t-end'].append(output[vid_idx]['segments'][:, 1])
                        results['label'].append(output[vid_idx]['labels'])
                        results['score'].append(output[vid_idx]['scores'])

                        pred_add = output[vid_idx]['segments'].detach().cpu().numpy() * output[vid_idx]['duration'] / val_loader.dataset.num_frames
                        gt_add =  np.array(output[vid_idx]['gt_time'])
                        score_add = output[vid_idx]['scores'].detach().cpu().numpy()
                        save_dict = {}
                        save_dict['gt_time'] = gt_add
                        save_dict['pred_time'] = pred_add
                        save_dict['score'] = score_add
                        result_save_obj.add_batch(save_dict)
        if (iter_idx != 0) and iter_idx % (print_freq) == 0:
            # measure elapsed time (sync all kernels)
            torch.cuda.synchronize()
            batch_time.update((time.time() - start) / print_freq)
            start = time.time()

            # print timing
            print('Test: [{0:05d}/{1:05d}]\t'
                  'Time {batch_time.val:.2f} ({batch_time.avg:.2f})'.format(
                  iter_idx, len(val_loader), batch_time=batch_time))

    # gather all stats and evaluate
    results['t-start'] = torch.cat(results['t-start']).numpy()
    results['t-end'] = torch.cat(results['t-end']).numpy()
    results['label'] = torch.cat(results['label']).numpy()
    results['score'] = torch.cat(results['score']).numpy()

    return None, acc_metrics_logger, result_save_obj_dict