<reponame>JiaMingLin/tsn-pytorch
import sys
sys.path.append("../")
import argparse
import os
import time
import shutil
import torch
import torchvision
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
from torch.nn.utils import clip_grad_norm

from dataset import TSNDataSet
from multi_task_model import TSN
from transforms import *
from opts import parser

from tensorboardX import SummaryWriter 
from math import ceil

# import warnings
# warnings.filterwarnings("ignore")


best_prec1 = 0

def main():
    global args, best_prec1, tb_writer_train, tb_writer_val, log_dir
    args = parser.parse_args()

    if args.dataset == 'ucf101':
        num_class = 101
    # elif args.dataset == 'hmdb51':
    #     num_class = 51
    # elif args.dataset == 'kinetics':
    #     num_class = 400
    else:
        raise ValueError('Unknown dataset '+args.dataset)

    log_dir = 'logs/{}'.format(args.save_path)
    tb_writer_train = SummaryWriter(os.path.join(log_dir, 'train'))
    tb_writer_val = SummaryWriter(os.path.join(log_dir, 'val'))

    model = TSN(num_class, args.num_segments,
                base_model=args.arch, new_length=args.new_length,
                consensus_type=args.consensus_type, dropout=args.dropout, partial_bn=not args.no_partialbn,
                modality=args.modality, resume=args.resume, project_mode = args.project_mode)

    crop_size = model.crop_size
    scale_size = model.scale_size
    input_mean = model.input_mean
    input_std = model.input_std
    policies = model.get_optim_policies()
    # train_augmentation = model.get_augmentation()

    device_ids = [int(id) for id in args.gpus.split(',')]
    model = torch.nn.DataParallel(model, device_ids=device_ids).cuda()

    cudnn.benchmark = True

    # Data loading code
    if args.modality != 'RGBDiff':
        normalize = GroupNormalize(input_mean, input_std)
    else:
        normalize = IdentityTransform()

    naming_pattern = "frame{:06d}.jpg"

    rgb_data_root_path = os.path.join(args.data_root_path, 'jpegs_256')
    tvl1_data_root_path = os.path.join(args.data_root_path, 'tvl1_flow')

    rgb_train_list = '../settings/ucf101/train_rgb_split1.txt'
    tvl1_train_list = '../settings/ucf101/train_flow_split1.txt'

    rgb_val_list = '../settings/ucf101/val_rgb_split1.txt'
    tvl1_val_list = '../settings/ucf101/val_flow_split1.txt'

    rgb_train_loader = torch.utils.data.DataLoader(
        TSNDataSet(rgb_data_root_path, rgb_train_list, num_segments=args.num_segments,
                   new_length=4,
                   modality="RGB",
                   image_tmpl=naming_pattern,
                   transform=torchvision.transforms.Compose([
                       torchvision.transforms.Compose([GroupMultiScaleCrop(224, [1, .875, .75, .66]),
                                                   GroupRandomHorizontalFlip(is_flow=False)]),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    tvl1_train_loader = torch.utils.data.DataLoader(
        TSNDataSet(tvl1_data_root_path, tvl1_train_list, num_segments=args.num_segments,
                   new_length=6,
                   modality="tvl1",
                   image_tmpl=naming_pattern,
                   transform=torchvision.transforms.Compose([
                       torchvision.transforms.Compose([GroupMultiScaleCrop(224, [1, .875, .75]),
                                                   GroupRandomHorizontalFlip(is_flow=True)]),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    rgb_val_loader = torch.utils.data.DataLoader(
        TSNDataSet(rgb_data_root_path, rgb_val_list, num_segments=args.num_segments,
                   new_length=4,
                   modality="RGB",
                   image_tmpl=naming_pattern,
                   random_shift=False,
                   transform=torchvision.transforms.Compose([
                       GroupScale(int(scale_size)),
                       GroupCenterCrop(crop_size),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    tvl1_val_loader = torch.utils.data.DataLoader(
        TSNDataSet(tvl1_data_root_path, tvl1_val_list, num_segments=args.num_segments,
                   new_length=6,
                   modality="tvl1",
                   image_tmpl=naming_pattern,
                   random_shift=False,
                   transform=torchvision.transforms.Compose([
                       GroupScale(int(scale_size)),
                       GroupCenterCrop(crop_size),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    train_loaders = [rgb_train_loader, tvl1_train_loader]
    val_loaders = [rgb_val_loader, tvl1_val_loader]

    # define loss function (criterion) and optimizer
    if args.loss_type == 'nll':
        criterion = torch.nn.CrossEntropyLoss().cuda()
    else:
        raise ValueError("Unknown loss type")

    for group in policies:
        print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(
            group['name'], len(group['params']), group['lr_mult'], group['decay_mult'])))

    optimizer = torch.optim.SGD(policies,
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    for epoch in range(args.start_epoch, args.epochs):
        adjust_learning_rate(optimizer, epoch, args.lr_steps)

        # train for one epoch
        avg_loss, avg_top1, avg_top5 = train(train_loaders, model, criterion, optimizer, epoch)
        tb_writer_train.add_scalar('Monitor/RGB Training Loss', avg_loss[0], epoch)
        tb_writer_train.add_scalar('Monitor/RGB Training Top1', avg_top1[0], epoch)
        tb_writer_train.add_scalar('Monitor/RGB Training Top5', avg_top5[0], epoch)

        tb_writer_train.add_scalar('Monitor/Flow Training Loss', avg_loss[1], epoch)
        tb_writer_train.add_scalar('Monitor/Flow Training Top1', avg_top1[1], epoch)
        tb_writer_train.add_scalar('Monitor/Flow Training Top5', avg_top5[1], epoch)

        # evaluate on validation set
        if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
            avg_loss, avg_top1, avg_top5 = validate(val_loaders, model, criterion)

            tb_writer_val.add_scalar('Monitor/RGB Validation Loss', avg_loss[0], epoch)
            tb_writer_val.add_scalar('Monitor/RGB Validation Top1', avg_top1[0], epoch)
            tb_writer_val.add_scalar('Monitor/RGB Validation Top5', avg_top5[0], epoch)

            tb_writer_val.add_scalar('Monitor/Flow Validation Loss', avg_loss[1], epoch)
            tb_writer_val.add_scalar('Monitor/Flow Validation Top1', avg_top1[1], epoch)
            tb_writer_val.add_scalar('Monitor/Flow Validation Top5', avg_top5[1], epoch)

            # remember best prec@1 and save checkpoint
            is_best = avg_top1[1] > best_prec1
            best_prec1 = max(avg_top1[1], best_prec1)
            save_checkpoint(log_dir, {
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.module.state_dict(),
                'best_prec1': best_prec1,
            }, is_best)


def train(train_loaders, model, criterion, optimizer, epoch):
    batch_time = [AverageMeter(), AverageMeter()]
    data_time = [AverageMeter(), AverageMeter()]
    losses = [AverageMeter(), AverageMeter()]
    top1 = [AverageMeter(), AverageMeter()]
    top5 = [AverageMeter(), AverageMeter()]

    model.module.partialBN(False)

    # switch to train mode
    model.train()

    end = time.time()

    train_loader_iter = [iter(loader) for loader in train_loaders]
    data_length = len(train_loaders)*(len(train_loaders[0]))

    for i in range(data_length):

        domain = i%len(train_loaders)
        data_time[domain].update(time.time() - end)

        loader_iter = train_loader_iter[domain]
        input_data, target = next(loader_iter)
        target = target.cuda()
        input_var = torch.autograd.Variable(input_data)
        target_var = torch.autograd.Variable(target)

        # compute output for a given domain
        output = model([input_var, domain])
        loss = criterion(output, target_var)

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1,5))
        losses[domain].update(loss.item(), input_data.size(0))
        top1[domain].update(prec1.item(), input_data.size(0))
        top5[domain].update(prec5.item(), input_data.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()

        loss.backward()

        if args.clip_gradient is not None:
            total_norm = clip_grad_norm(model.parameters(), args.clip_gradient)
            #if total_norm > args.clip_gradient:
            #    print("clipping gradient: {} with coef {}".format(total_norm, args.clip_gradient / total_norm))

        optimizer.step()

        # measure elapsed time
        batch_time[domain].update(time.time() - end)
        end = time.time()

        if ceil(i/len(train_loaders)) % args.print_freq == 0 and i > 0:
            print(('Domain: {0}, Epoch: [{1}][{2}/{3}], lr: {lr:.5f}\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                   "RGB" if domain == 0 else "tvl1", 
                   epoch, ceil(i/len(train_loaders)), len(train_loaders[domain]), batch_time=batch_time[domain],
                   data_time=data_time[domain], loss=losses[domain], top1=top1[domain], top5=top5[domain], lr=optimizer.param_groups[-1]['lr'])))
    return [loss.avg for loss in losses], [t1.avg for t1 in top1], [t5.avg for t5 in top5]

def validate(val_loaders, model, criterion, logger=None):
    batch_time = [AverageMeter(), AverageMeter()]
    losses = [AverageMeter(), AverageMeter()]
    top1 = [AverageMeter(), AverageMeter()]
    top5 = [AverageMeter(), AverageMeter()]

    # switch to evaluate mode
    model.eval()

    val_loader_iters = [iter(loader) for loader in val_loaders]

    data_length = len(val_loaders) * len(val_loaders[0])

    end = time.time()
    with torch.no_grad():

        for domain, loader in enumerate(val_loaders):
            for i, (input_data, target) in enumerate(loader):
                target = target.cuda()
                input_var = torch.autograd.Variable(input_data)
                target_var = torch.autograd.Variable(target)

                # compute output
                output = model([input_var, domain])
                loss = criterion(output, target_var)

                # measure accuracy and record loss
                prec1, prec5 = accuracy(output.data, target, topk=(1,5))

                losses[domain].update(loss.data.item(), input_data.size(0))
                top1[domain].update(prec1.item(), input_data.size(0))
                top5[domain].update(prec5.item(), input_data.size(0))

                # measure elapsed time
                batch_time[domain].update(time.time() - end)
                end = time.time()

                if i % args.print_freq == 0:
                    print(('Domain: {0},Test: [{1}/{2}]\t'
                       'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                       'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                       'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                       'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                            "RGB" if domain == 0 else "tvl1",
                            i, len(val_loaders[domain]), batch_time=batch_time[domain], loss=losses[domain],
                            top1=top1[domain], top5=top5[domain])))

            print(('Domain: {0}, Testing Results: Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Loss {loss.avg:.5f}'
                    .format("RGB" if domain == 0 else "tvl1", top1=top1[domain], top5=top5[domain], loss=losses[domain])))

    return [loss.avg for loss in losses], [t1.avg for t1 in top1], [t5.avg for t5 in top5]


def save_checkpoint(log_dir, state, is_best, filename='checkpoint.pth.tar'):
    
    filename = '_'.join(
        (
            'epoch_{}'.format(state['epoch']), 
            '{:0.2f}'.format(state['best_prec1']),
            state['arch'], 
            args.modality.lower(), 
            filename)
        )
    save_path = os.path.join(log_dir, filename)
    torch.save(state, save_path)
    if is_best:
        best_name = '_'.join(
            (
                'best',
                'epoch_{}'.format(state['epoch']), 
                '{:0.2f}'.format(state['best_prec1']),
                args.modality.lower(), 
                '.pth.tar')
            )
        best_path = os.path.join(log_dir, best_name)
        shutil.copyfile(save_path, best_path)


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

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


def adjust_learning_rate(optimizer, epoch, lr_steps):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    decay = 0.1 ** (sum(epoch >= np.array(lr_steps)))
    lr = args.lr * decay
    decay = args.weight_decay
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr * param_group['lr_mult']
        param_group['weight_decay'] = decay * param_group['decay_mult']


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].contiguous().view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


if __name__ == '__main__':
    ## Log directory
    main()

import sys
sys.path.append("../")
import argparse
import os
import time
import shutil
import torch
import torchvision
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
from torch.nn.utils import clip_grad_norm

from dataset import TSNDataSet
from multi_task_model import TSN
from transforms import *
from opts import parser

from tensorboardX import SummaryWriter 
from math import ceil

# import warnings
# warnings.filterwarnings("ignore")


best_prec1 = 0

def main():
    global args, best_prec1, tb_writer_train, tb_writer_val, log_dir
    args = parser.parse_args()

    if args.dataset == 'ucf101':
        num_class = 101
    # elif args.dataset == 'hmdb51':
    #     num_class = 51
    # elif args.dataset == 'kinetics':
    #     num_class = 400
    else:
        raise ValueError('Unknown dataset '+args.dataset)

    log_dir = 'logs/{}'.format(args.save_path)
    tb_writer_train = SummaryWriter(os.path.join(log_dir, 'train'))
    tb_writer_val = SummaryWriter(os.path.join(log_dir, 'val'))

    model = TSN(num_class, args.num_segments,
                base_model=args.arch, new_length=args.new_length,
                consensus_type=args.consensus_type, dropout=args.dropout, partial_bn=not args.no_partialbn,
                modality=args.modality, resume=args.resume, project_mode = args.project_mode)

    crop_size = model.crop_size
    scale_size = model.scale_size
    input_mean = model.input_mean
    input_std = model.input_std
    policies = model.get_optim_policies()
    # train_augmentation = model.get_augmentation()

    device_ids = [int(id) for id in args.gpus.split(',')]
    model = torch.nn.DataParallel(model, device_ids=device_ids).cuda()

    cudnn.benchmark = True

    # Data loading code
    if args.modality != 'RGBDiff':
        normalize = GroupNormalize(input_mean, input_std)
    else:
        normalize = IdentityTransform()

    naming_pattern = "frame{:06d}.jpg"

    rgb_data_root_path = os.path.join(args.data_root_path, 'jpegs_256')
    tvl1_data_root_path = os.path.join(args.data_root_path, 'tvl1_flow')

    rgb_train_list = '../settings/ucf101/train_rgb_split1.txt'
    tvl1_train_list = '../settings/ucf101/train_flow_split1.txt'

    rgb_val_list = '../settings/ucf101/val_rgb_split1.txt'
    tvl1_val_list = '../settings/ucf101/val_flow_split1.txt'

    rgb_train_loader = torch.utils.data.DataLoader(
        TSNDataSet(rgb_data_root_path, rgb_train_list, num_segments=args.num_segments,
                   new_length=4,
                   modality="RGB",
                   image_tmpl=naming_pattern,
                   transform=torchvision.transforms.Compose([
                       torchvision.transforms.Compose([GroupMultiScaleCrop(224, [1, .875, .75, .66]),
                                                   GroupRandomHorizontalFlip(is_flow=False)]),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    tvl1_train_loader = torch.utils.data.DataLoader(
        TSNDataSet(tvl1_data_root_path, tvl1_train_list, num_segments=args.num_segments,
                   new_length=6,
                   modality="tvl1",
                   image_tmpl=naming_pattern,
                   transform=torchvision.transforms.Compose([
                       torchvision.transforms.Compose([GroupMultiScaleCrop(224, [1, .875, .75]),
                                                   GroupRandomHorizontalFlip(is_flow=True)]),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    rgb_val_loader = torch.utils.data.DataLoader(
        TSNDataSet(rgb_data_root_path, rgb_val_list, num_segments=args.num_segments,
                   new_length=4,
                   modality="RGB",
                   image_tmpl=naming_pattern,
                   random_shift=False,
                   transform=torchvision.transforms.Compose([
                       GroupScale(int(scale_size)),
                       GroupCenterCrop(crop_size),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    tvl1_val_loader = torch.utils.data.DataLoader(
        TSNDataSet(tvl1_data_root_path, tvl1_val_list, num_segments=args.num_segments,
                   new_length=6,
                   modality="tvl1",
                   image_tmpl=naming_pattern,
                   random_shift=False,
                   transform=torchvision.transforms.Compose([
                       GroupScale(int(scale_size)),
                       GroupCenterCrop(crop_size),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    train_loaders = [rgb_train_loader, tvl1_train_loader]
    val_loaders = [rgb_val_loader, tvl1_val_loader]

    # define loss function (criterion) and optimizer
    if args.loss_type == 'nll':
        criterion = torch.nn.CrossEntropyLoss().cuda()
    else:
        raise ValueError("Unknown loss type")

    for group in policies:
        print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(
            group['name'], len(group['params']), group['lr_mult'], group['decay_mult'])))

    optimizer = torch.optim.SGD(policies,
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    for epoch in range(args.start_epoch, args.epochs):
        adjust_learning_rate(optimizer, epoch, args.lr_steps)

        # train for one epoch
        avg_loss, avg_top1, avg_top5 = train(train_loaders, model, criterion, optimizer, epoch)
        tb_writer_train.add_scalar('Monitor/RGB Training Loss', avg_loss[0], epoch)
        tb_writer_train.add_scalar('Monitor/RGB Training Top1', avg_top1[0], epoch)
        tb_writer_train.add_scalar('Monitor/RGB Training Top5', avg_top5[0], epoch)

        tb_writer_train.add_scalar('Monitor/Flow Training Loss', avg_loss[1], epoch)
        tb_writer_train.add_scalar('Monitor/Flow Training Top1', avg_top1[1], epoch)
        tb_writer_train.add_scalar('Monitor/Flow Training Top5', avg_top5[1], epoch)

        # evaluate on validation set
        if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
            avg_loss, avg_top1, avg_top5 = validate(val_loaders, model, criterion)

            tb_writer_val.add_scalar('Monitor/RGB Validation Loss', avg_loss[0], epoch)
            tb_writer_val.add_scalar('Monitor/RGB Validation Top1', avg_top1[0], epoch)
            tb_writer_val.add_scalar('Monitor/RGB Validation Top5', avg_top5[0], epoch)

            tb_writer_val.add_scalar('Monitor/Flow Validation Loss', avg_loss[1], epoch)
            tb_writer_val.add_scalar('Monitor/Flow Validation Top1', avg_top1[1], epoch)
            tb_writer_val.add_scalar('Monitor/Flow Validation Top5', avg_top5[1], epoch)

            # remember best prec@1 and save checkpoint
            is_best = avg_top1[1] > best_prec1
            best_prec1 = max(avg_top1[1], best_prec1)
            save_checkpoint(log_dir, {
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.module.state_dict(),
                'best_prec1': best_prec1,
            }, is_best)


def train(train_loaders, model, criterion, optimizer, epoch):
    batch_time = [AverageMeter(), AverageMeter()]
    data_time = [AverageMeter(), AverageMeter()]
    losses = [AverageMeter(), AverageMeter()]
    top1 = [AverageMeter(), AverageMeter()]
    top5 = [AverageMeter(), AverageMeter()]

    model.module.partialBN(False)

    # switch to train mode
    model.train()

    end = time.time()

    train_loader_iter = [iter(loader) for loader in train_loaders]
    data_length = len(train_loaders)*(len(train_loaders[0]))

    for i in range(data_length):

        domain = i%len(train_loaders)
        data_time[domain].update(time.time() - end)

        loader_iter = train_loader_iter[domain]
        input_data, target = next(loader_iter)
        target = target.cuda()
        input_var = torch.autograd.Variable(input_data)
        target_var = torch.autograd.Variable(target)

        # compute output for a given domain
        output = model([input_var, domain])
        loss = criterion(output, target_var)

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1,5))
        losses[domain].update(loss.item(), input_data.size(0))
        top1[domain].update(prec1.item(), input_data.size(0))
        top5[domain].update(prec5.item(), input_data.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()

        loss.backward()

        if args.clip_gradient is not None:
            total_norm = clip_grad_norm(model.parameters(), args.clip_gradient)
            #if total_norm > args.clip_gradient:
            #    print("clipping gradient: {} with coef {}".format(total_norm, args.clip_gradient / total_norm))

        optimizer.step()

        # measure elapsed time
        batch_time[domain].update(time.time() - end)
        end = time.time()

        if ceil(i/len(train_loaders)) % args.print_freq == 0 and i > 0:
            print(('Domain: {0}, Epoch: [{1}][{2}/{3}], lr: {lr:.5f}\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                   "RGB" if domain == 0 else "tvl1", 
                   epoch, ceil(i/len(train_loaders)), len(train_loaders[domain]), batch_time=batch_time[domain],
                   data_time=data_time[domain], loss=losses[domain], top1=top1[domain], top5=top5[domain], lr=optimizer.param_groups[-1]['lr'])))
    return [loss.avg for loss in losses], [t1.avg for t1 in top1], [t5.avg for t5 in top5]

def validate(val_loaders, model, criterion, logger=None):
    batch_time = [AverageMeter(), AverageMeter()]
    losses = [AverageMeter(), AverageMeter()]
    top1 = [AverageMeter(), AverageMeter()]
    top5 = [AverageMeter(), AverageMeter()]

    # switch to evaluate mode
    model.eval()

    val_loader_iters = [iter(loader) for loader in val_loaders]

    data_length = len(val_loaders) * len(val_loaders[0])

    end = time.time()
    with torch.no_grad():

        for domain, loader in enumerate(val_loaders):
            for i, (input_data, target) in enumerate(loader):
                target = target.cuda()
                input_var = torch.autograd.Variable(input_data)
                target_var = torch.autograd.Variable(target)

                # compute output
                output = model([input_var, domain])
                loss = criterion(output, target_var)

                # measure accuracy and record loss
                prec1, prec5 = accuracy(output.data, target, topk=(1,5))

                losses[domain].update(loss.data.item(), input_data.size(0))
                top1[domain].update(prec1.item(), input_data.size(0))
                top5[domain].update(prec5.item(), input_data.size(0))

                # measure elapsed time
                batch_time[domain].update(time.time() - end)
                end = time.time()

                if i % args.print_freq == 0:
                    print(('Domain: {0},Test: [{1}/{2}]\t'
                       'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                       'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                       'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                       'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                            "RGB" if domain == 0 else "tvl1",
                            i, len(val_loaders[domain]), batch_time=batch_time[domain], loss=losses[domain],
                            top1=top1[domain], top5=top5[domain])))

            print(('Domain: {0}, Testing Results: Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Loss {loss.avg:.5f}'
                    .format("RGB" if domain == 0 else "tvl1", top1=top1[domain], top5=top5[domain], loss=losses[domain])))

    return [loss.avg for loss in losses], [t1.avg for t1 in top1], [t5.avg for t5 in top5]


def save_checkpoint(log_dir, state, is_best, filename='checkpoint.pth.tar'):
    
    filename = '_'.join(
        (
            'epoch_{}'.format(state['epoch']), 
            '{:0.2f}'.format(state['best_prec1']),
            state['arch'], 
            args.modality.lower(), 
            filename)
        )
    save_path = os.path.join(log_dir, filename)
    torch.save(state, save_path)
    if is_best:
        best_name = '_'.join(
            (
                'best',
                'epoch_{}'.format(state['epoch']), 
                '{:0.2f}'.format(state['best_prec1']),
                args.modality.lower(), 
                '.pth.tar')
            )
        best_path = os.path.join(log_dir, best_name)
        shutil.copyfile(save_path, best_path)


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

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


def adjust_learning_rate(optimizer, epoch, lr_steps):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    decay = 0.1 ** (sum(epoch >= np.array(lr_steps)))
    lr = args.lr * decay
    decay = args.weight_decay
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr * param_group['lr_mult']
        param_group['weight_decay'] = decay * param_group['decay_mult']


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].contiguous().view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


if __name__ == '__main__':
    ## Log directory
    main()
# import warnings # warnings.filterwarnings("ignore") # elif args.dataset == 'hmdb51': #     num_class = 51 # elif args.dataset == 'kinetics': #     num_class = 400 # train_augmentation = model.get_augmentation() # Data loading code # define loss function (criterion) and optimizer # train for one epoch # evaluate on validation set # remember best prec@1 and save checkpoint # switch to train mode # compute output for a given domain # measure accuracy and record loss # compute gradient and do SGD step #if total_norm > args.clip_gradient: #    print("clipping gradient: {} with coef {}".format(total_norm, args.clip_gradient / total_norm)) # measure elapsed time # switch to evaluate mode # compute output # measure accuracy and record loss # measure elapsed time Computes and stores the average and current value Sets the learning rate to the initial LR decayed by 10 every 30 epochs Computes the precision@k for the specified values of k ## Log directory