engine.py

# Modified from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# ------------------------------------------------------------------------
"""
Modified from RTD-Net (https://github.com/MCG-NJU/RTD-Action)

PointTAD Training and Inference functions.

"""

import json
import math
import sys
from typing import Iterable

import torch
from termcolor import colored

import util.misc as utils
from datasets.evaluate import Evaluator, cal_multi_map


def load_json(file):
    with open(file) as json_file:
        data = json.load(json_file)
        return data


def train_one_epoch(model: torch.nn.Module,
                    criterion: torch.nn.Module,
                    data_loader: Iterable,
                    optimizer: torch.optim.Optimizer,
                    device: torch.device,
                    epoch: int,
                    args,
                    model_ema=None,
                    model_for_ema=None):
    model.train()
    criterion.train()
    metric_logger = utils.MetricLogger(delimiter='  ')
    metric_logger.add_meter(
        'lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    
    # metric_logger.add_meter(
    #     'class_error', utils.SmoothedValue(window_size=1,
    #                                         fmt='{value:.2f}'))
    header = 'Epoch: [{}]'.format(epoch)
    print_freq = args.print_freq

    max_norm = args.clip_max_norm

    for vid_name_list, locations, x, targets, num_frames, base \
        in metric_logger.log_every(data_loader, print_freq, header):

        x = x.to(device)
        locations = locations.to(device)
        if args.model_name=='pointtad':
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        bs = x.shape[0]
        outputs = model(x, locations=locations)
        loss_dict = criterion(outputs, targets, epoch)
        weight_dict = criterion.weight_dict
        losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys()
                     if k in weight_dict)
        # n_parameters = sum(p.numel() for p in model.parameters())
        # losses += 0 * n_parameters
        
        loss_dict_reduced = utils.reduce_dict(loss_dict)
        loss_dict_reduced_unscaled = {
            f'{k}_unscaled': v
            for k, v in loss_dict_reduced.items()}
        loss_dict_reduced_scaled = {
            k: v * weight_dict[k]
            for k, v in loss_dict_reduced.items() if k in weight_dict}
        losses_reduced_scaled = sum(loss_dict_reduced_scaled.values())
        loss_value = losses_reduced_scaled.item()

        if not math.isfinite(loss_value):
            print('Loss is {}, stopping training'.format(loss_value))
            print(loss_dict_reduced)
            sys.exit(1)

        optimizer.zero_grad()
        losses.backward()
        if max_norm > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        optimizer.step()
        if (
            model_ema is not None
            and model_for_ema is not None
            and epoch >= args.model_ema_start_epoch
        ):
            model_ema.update_parameters(model_for_ema)

        metric_logger.update(loss=loss_value,
                             **loss_dict_reduced_scaled,
                             **loss_dict_reduced_unscaled)
        if 'class_error' in loss_dict_reduced:
            metric_logger.update(class_error=loss_dict_reduced['class_error'])
        metric_logger.update(lr=optimizer.param_groups[0]['lr'])

    metric_logger.synchronize_between_processes()
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}, loss_value # loss_dict


@torch.no_grad()
def evaluate(model, criterion, postprocessors, data_loader, device, args):
    print(colored('evaluate', 'red'))
    model.eval()
    criterion.eval()

    metric_logger = utils.MetricLogger(delimiter='  ')
    header = 'Test:'
    val_loss = 0.0
    num_iter = 0
    evaluator = Evaluator()
    video_pool = list(load_json(args.annotation_path).keys())
    video_pool.sort()
    video_dict = {i: video_pool[i] for i in range(len(video_pool))}
    for vid_name_list, locations, x, targets, num_frames, base in metric_logger.log_every(
            data_loader, args.print_freq, header):
        x = x.to(device)
        locations = locations.to(device)
        if args.model_name=='pointtad':
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        num_iter+=1
        outputs = model(x, locations)
        loss_dict = criterion(outputs, targets)
        weight_dict = criterion.weight_dict
        losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys()
                     if k in weight_dict)
        val_loss += losses.item()
        
        # print reduced values
        loss_dict_reduced = utils.reduce_dict(loss_dict)
        loss_dict_reduced_scaled = {
            k: v * weight_dict[k]
            for k, v in loss_dict_reduced.items() if k in weight_dict
        }
        loss_dict_reduced_unscaled = {
            f'{k}_unscaled': v
            for k, v in loss_dict_reduced.items()
        }
        metric_logger.update(loss=sum(loss_dict_reduced_scaled.values()),
                             **loss_dict_reduced_scaled,
                             **loss_dict_reduced_unscaled)
        if 'class_error' in loss_dict_reduced:
            metric_logger.update(class_error=loss_dict_reduced['class_error'])

        results, dense_results = postprocessors['results'](outputs, num_frames, base, targets)
        for target, output, dense_res, base_loc in zip(targets, results, dense_results, base):
            vid = video_dict[target['video_id'].item()]
            # dense_gt = target['dense_gt']
            if args.model_name=='freetad':
                evaluator.update(vid, output, base_loc, target)
            else:
                evaluator.update(vid, output, base_loc)
            
            if args.dense_result:
                torch.save(dense_res, f'dense_results/{vid}_{base_loc}_dense')
                # all_label.append(dense_gt.cpu())
                # all_res.append(dense_res.cpu())

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    evaluator.synchronize_between_processes()
    print('Averaged stats:', metric_logger)
    # if args.dense_result:
    #     seg_map, _ = cal_multi_map(torch.cat(all_res), torch.cat(all_label))
    #     print('Segmentation map:', seg_map)

    return evaluator, val_loss/num_iter