engine_for_vsr.py

# --------------------------------------------------------
# Based on BEiT, timm, DINO and DeiT code bases
# https://github.com/microsoft/unilm/tree/master/beit
# https://github.com/rwightman/pytorch-image-models/tree/master/timm
# https://github.com/facebookresearch/deit
# https://github.com/facebookresearch/dino
# --------------------------------------------------------'
import math
import os
import sys
from multiprocessing import Pool
from typing import Iterable, Optional

import numpy as np
import torch
from scipy.special import softmax
from timm.data import Mixup
from timm.utils import ModelEma, accuracy

import utils

from torch.nn.utils.rnn import pad_sequence
from decoders import ctc_greedy_decode, compute_CTC_prob
from metrics import compute_error_ch, compute_error_word

def train_class_batch(model, samples, target, src_len, tgt_len, criterion):
    Alpha = 0.2
    inputLenBatch, outputBatch = model(samples, target, src_len, tgt_len)

    tgt_list = [target[i][:tgt_len[i]-1] for i in range(len(tgt_len))]
    tgt_out_batch = pad_sequence(tgt_list, batch_first=True)
    targetMask = torch.zeros_like(tgt_out_batch, device=samples.device)
    targetMask[(torch.arange(targetMask.shape[0]), tgt_len.long() - 2)] = 1
    targetMask = (1 - targetMask.flip([-1]).cumsum(-1).flip([-1])).bool()
    concatTargetoutBatch = tgt_out_batch[~targetMask]

    with torch.backends.cudnn.flags(enabled=False):
        ctcloss = model.CTCloss(outputBatch[0], concatTargetoutBatch, inputLenBatch, tgt_len-1)
        celoss = model.CEloss(outputBatch[1], tgt_out_batch.long())
        loss = Alpha * ctcloss + (1 - Alpha) * celoss
    if loss is torch.nan:
        print(loss)
        # breakpoint()

    predictionBatch, predictionLenBatch = ctc_greedy_decode(outputBatch[0].detach(), inputLenBatch, CHAR_TO_INDEX['<EOS>'])
    c_edits, c_count = compute_error_ch(predictionBatch, concatTargetoutBatch, predictionLenBatch, tgt_len-1)
    cer = c_edits / c_count
    w_edits, w_count = compute_error_word(predictionBatch, concatTargetoutBatch, predictionLenBatch, tgt_len-1, CHAR_TO_INDEX[' '])
    wer = w_edits / w_count
    return loss, cer, wer


def get_loss_scale_for_deepspeed(model):
    optimizer = model.optimizer
    return optimizer.loss_scale if hasattr(
        optimizer, "loss_scale") else optimizer.cur_scale


CHAR_TO_INDEX= {" ": 1, "'": 22, "1": 30, "0": 29, "3": 37, "2": 32, "5": 34, "4": 38, "7": 36, "6": 35, "9": 31, "8": 33, "A": 5, "C": 17,
                "B": 20, "E": 2, "D": 12, "G": 16, "F": 19, "I": 6, "H": 9, "K": 24, "J": 25, "M": 18, "L": 11, "O": 4, "N": 7, "Q": 27,
                "P": 21, "S": 8, "R": 10, "U": 13, "T": 3, "W": 15, "V": 23, "Y": 14, "X": 26, "Z": 28, "<EOS>": 39}

def text2idx(sentence):
    ans = [CHAR_TO_INDEX[c] for c in sentence.upper()]
    ans.append(CHAR_TO_INDEX['<EOS>'])
    return ans


def train_one_epoch(model: torch.nn.Module,
                    criterion: torch.nn.Module,
                    data_loader: Iterable,
                    optimizer: torch.optim.Optimizer,
                    device: torch.device,
                    epoch: int,
                    loss_scaler,
                    max_norm: float = 0,
                    model_ema: Optional[ModelEma] = None,
                    mixup_fn: Optional[Mixup] = None,
                    log_writer=None,
                    start_steps=None,
                    lr_schedule_values=None,
                    wd_schedule_values=None,
                    num_training_steps_per_epoch=None,
                    update_freq=None):
    model.train(True)
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter(
        'lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    metric_logger.add_meter(
        'min_lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    header = 'Epoch: [{}]'.format(epoch)
    print_freq = 20

    if loss_scaler is None:
        model.zero_grad()
        model.micro_steps = 0
    else:
        optimizer.zero_grad()

    for data_iter_step, (samples, targets, src_len, tgt_len) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
        step = data_iter_step // update_freq
        if step >= num_training_steps_per_epoch:
            continue
        it = start_steps + step  # global training iteration
        # Update LR & WD for the first acc
        if lr_schedule_values is not None or wd_schedule_values is not None and data_iter_step % update_freq == 0:
            for i, param_group in enumerate(optimizer.param_groups):
                if lr_schedule_values is not None:
                    param_group["lr"] = lr_schedule_values[it] * param_group[
                        "lr_scale"]
                if wd_schedule_values is not None and param_group[
                        "weight_decay"] > 0:
                    param_group["weight_decay"] = wd_schedule_values[it]

        samples = samples.to(device, non_blocking=True)
        # tgt = [torch.tensor(text2idx(s)) for s in targets]
        # tgt = pad_sequence(tgt, batch_first = True)
        targets = targets.to(device, non_blocking=True)

        if mixup_fn is not None:
            # mixup handle 3th & 4th dimension
            B, C, T, H, W = samples.shape
            samples = samples.view(B, C * T, H, W)
            samples, targets = mixup_fn(samples, targets)
            samples = samples.view(B, C, T, H, W)

        if loss_scaler is None:
            samples = samples.half()
            loss, cer, wer = train_class_batch(model, samples, targets, src_len, tgt_len, criterion)
        else:
            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                loss, cer, wer = train_class_batch(model, samples, targets, src_len, tgt_len, criterion)

        loss_value = loss.item()

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

        if loss_scaler is None:
            loss /= update_freq
            model.backward(loss)
            grad_norm = model.get_global_grad_norm()

            model.step()

            if (data_iter_step + 1) % update_freq == 0:
                # Deepspeed will call step() & model.zero_grad() automatic
                if model_ema is not None:
                    model_ema.update(model)
            loss_scale_value = get_loss_scale_for_deepspeed(model)
        else:
            # this attribute is added by timm on one optimizer (adahessian)
            is_second_order = hasattr(
                optimizer, 'is_second_order') and optimizer.is_second_order
            loss /= update_freq
            grad_norm = loss_scaler(
                loss,
                optimizer,
                clip_grad=max_norm,
                parameters=model.parameters(),
                create_graph=is_second_order,
                update_grad=(data_iter_step + 1) % update_freq == 0)
            if (data_iter_step + 1) % update_freq == 0:
                optimizer.zero_grad()
                if model_ema is not None:
                    model_ema.update(model)
            loss_scale_value = loss_scaler.state_dict()["scale"]

        torch.cuda.synchronize()

        # if mixup_fn is None:
        #     class_acc = (output.max(-1)[-1] == targets).float().mean()
        # else:
        #     class_acc = None
        metric_logger.update(loss=loss_value)
        metric_logger.update(cer=cer)
        metric_logger.update(wer=wer)
        metric_logger.update(loss_scale=loss_scale_value)
        min_lr = 10.
        max_lr = 0.
        for group in optimizer.param_groups:
            min_lr = min(min_lr, group["lr"])
            max_lr = max(max_lr, group["lr"])

        metric_logger.update(lr=max_lr)
        metric_logger.update(min_lr=min_lr)
        weight_decay_value = None
        for group in optimizer.param_groups:
            if group["weight_decay"] > 0:
                weight_decay_value = group["weight_decay"]
        metric_logger.update(weight_decay=weight_decay_value)
        metric_logger.update(grad_norm=grad_norm)

        if log_writer is not None:
            log_writer.update(loss=loss_value, head="loss")
            log_writer.update(cer=cer, head="loss")
            log_writer.update(wer=wer, head="loss")
            log_writer.update(loss_scale=loss_scale_value, head="opt")
            log_writer.update(lr=max_lr, head="opt")
            log_writer.update(min_lr=min_lr, head="opt")
            log_writer.update(weight_decay=weight_decay_value, head="opt")
            log_writer.update(grad_norm=grad_norm, head="opt")

            log_writer.set_step()

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}


@torch.no_grad()
def validation_one_epoch(data_loader, model, device):
    criterion = torch.nn.CrossEntropyLoss()

    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Val:'

    # switch to evaluation mode
    model.eval()

    for batch in metric_logger.log_every(data_loader, 10, header):
        src = batch[0]
        tgt = batch[1]
        src_len = batch[2]
        tgt_len = batch[3]
        src = src.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)

        # compute output
        with torch.cuda.amp.autocast():
            predictionBatch, predictionLenBatch = model.inference(src, tgt, src_len, tgt_len, criterion)

        batch_size = src.shape[0]
        tgt_list = [target[i][:tgt_len[i]-1] for i in range(len(tgt_len))]
        tgt_out_batch = pad_sequence(tgt_list, batch_first=True)
        targetMask = torch.zeros_like(tgt_out_batch, device=device)
        targetMask[(torch.arange(targetMask.shape[0]), tgt_len.long() - 2)] = 1
        targetMask = (1 - targetMask.flip([-1]).cumsum(-1).flip([-1])).bool()
        concatTargetoutBatch = tgt_out_batch[~targetMask]

        c_edits, c_count = compute_error_ch(predictionBatch, concatTargetoutBatch, predictionLenBatch, tgt_len-1)
        w_edits, w_count = compute_error_word(predictionBatch, concatTargetoutBatch, predictionLenBatch, tgt_len-1, CHAR_TO_INDEX[' '])

        metric_logger.update(loss=0)
        metric_logger.meters['cer'].update(c_edits, n=c_count)
        metric_logger.meters['wer'].update(w_edits, n=w_count)
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print(
        '* CER {top1.global_avg:.3f} WER {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
        .format(
            top1=metric_logger.cer,
            top5=metric_logger.wer,
            losses=metric_logger.loss))

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


@torch.no_grad()
def final_test(data_loader, model, device, file):
    evalCER = 0
    evalWER = 0
    evalCCount = 0
    evalWCount = 0
    criterion = torch.nn.CrossEntropyLoss()

    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Test:'

    # switch to evaluation mode
    model.eval()
    final_result = []

    for batch in metric_logger.log_every(data_loader, 10, header):
        src = batch[0]
        tgt = batch[1]
        src_len = batch[2]
        tgt_len = batch[3]
        src = src.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)

        # compute output
        with torch.cuda.amp.autocast():
            predictionBatch, predictionLenBatch = model.inference(src, tgt, src_len, tgt_len, criterion)

        batch_size = src.shape[0]
        tgt_list = [target[i][:tgt_len[i]-1] for i in range(len(tgt_len))]
        tgt_out_batch = pad_sequence(tgt_list, batch_first=True)
        targetMask = torch.zeros_like(tgt_out_batch, device=device)
        targetMask[(torch.arange(targetMask.shape[0]), tgt_len.long() - 2)] = 1
        targetMask = (1 - targetMask.flip([-1]).cumsum(-1).flip([-1])).bool()
        concatTargetoutBatch = tgt_out_batch[~targetMask]

        c_edits, c_count = compute_error_ch(predictionBatch, concatTargetoutBatch, predictionLenBatch, tgt_len-1)
        w_edits, w_count = compute_error_word(predictionBatch, concatTargetoutBatch, predictionLenBatch, tgt_len-1, CHAR_TO_INDEX[' '])
        
        evalCER += c_edits
        evalCCount += c_count
        evalWER += w_edits
        evalWCount += w_count

        metric_logger.update(loss=0)
        metric_logger.meters['cer'].update(c_edits, n=c_count)
        metric_logger.meters['wer'].update(w_edits, n=w_count)

    if not os.path.exists(file):
        os.mknod(file)
    with open(file, 'w') as f:
        f.write("{}, {}\n".format(evalCER/evalCCount if evalCCount > 0 else 1, evalWER/evalWCount if evalWCount > 0 else 1))
        for line in final_result:
            f.write(line)
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print(
        '* CER {top1.global_avg:.3f} WER {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
        .format(
            top1=metric_logger.cer,
            top5=metric_logger.wer,
            losses=metric_logger.loss))

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


def merge(eval_path, num_tasks, method='prob'):
    assert method in ['prob', 'score']
    dict_feats = {}
    dict_label = {}
    dict_pos = {}
    print("Reading individual output files")

    for x in range(num_tasks):
        file = os.path.join(eval_path, str(x) + '.txt')
        lines = open(file, 'r').readlines()[1:]
        for line in lines:
            line = line.strip()
            name = line.split('[')[0]
            label = line.split(']')[1].split(' ')[1]
            chunk_nb = line.split(']')[1].split(' ')[2]
            split_nb = line.split(']')[1].split(' ')[3]
            data = np.fromstring(
                line.split('[')[1].split(']')[0], dtype=float, sep=',')
            if name not in dict_feats:
                dict_feats[name] = []
                dict_label[name] = 0
                dict_pos[name] = []
            if chunk_nb + split_nb in dict_pos[name]:
                continue
            if method == 'prob':
                dict_feats[name].append(softmax(data))
            else:
                dict_feats[name].append(data)
            dict_pos[name].append(chunk_nb + split_nb)
            dict_label[name] = label
    print("Computing final results")

    input_lst = []
    for i, item in enumerate(dict_feats):
        input_lst.append([i, item, dict_feats[item], dict_label[item]])
    p = Pool(64)
    # [pred, top1, top5, label]
    ans = p.map(compute_video, input_lst)
    top1 = [x[1] for x in ans]
    top5 = [x[2] for x in ans]
    label = [x[3] for x in ans]
    final_top1, final_top5 = np.mean(top1), np.mean(top5)

    return final_top1 * 100, final_top5 * 100


def compute_video(lst):
    i, video_id, data, label = lst
    feat = [x for x in data]
    feat = np.mean(feat, axis=0)
    pred = np.argmax(feat)
    top1 = (int(pred) == int(label)) * 1.0
    top5 = (int(label) in np.argsort(-feat)[:5]) * 1.0
    return [pred, top1, top5, int(label)]


def collate_func(batch):
    src = pad_sequence([data[0].permute(1,0,2,3) for data in batch], batch_first=True).transpose(1,2)
    # B x 3 x T x H x W
    tgt = pad_sequence([torch.tensor(text2idx(data[1])) for data in batch], batch_first=True)
    # B x T
    src_len = torch.tensor([data[2] for data in batch])
    tgt_len = torch.tensor([data[3]+1 for data in batch])
    return src, tgt, src_len, tgt_len