train.py

# Copyright (C) 2025-present Meta Platforms, Inc. and affiliates. All rights reserved.
# Licensed under CC BY-NC 4.0 (non-commercial use only).


#!/usr/bin/env python3
import os, sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__)))

if 'META_INTERNAL' in os.environ.keys() and os.environ['META_INTERNAL'] == "False":
    generate_html = None
    from dust3r.dummy_io import *
else:
    from meta_internal.io import *
    from meta_internal.html_gen.run_model_doctor import generate_html

import argparse
import datetime
import json
import numpy as np
import time
import math
from collections import defaultdict
from pathlib import Path
from typing import Sized
import imageio


import torch
import torch.backends.cudnn as cudnn
from torch.utils.tensorboard import SummaryWriter
torch.backends.cuda.matmul.allow_tf32 = True  # for gpu >= Ampere and pytorch >= 1.12

from dust3r.model import AsymmetricCroCo3DStereo, AsymmetricCroCo3DStereoMultiView, inf  # noqa: F401, needed when loading the model
from dust3r.datasets import get_data_loader  # noqa
from dust3r.losses import *  # noqa: F401, needed when loading the model
from dust3r.inference import loss_of_one_batch  # noqa
from dust3r.pcd_render import pcd_render, save_video_combined
from dust3r.gs import gs_render
from dust3r.utils.geometry import inv, geotrf

import dust3r.utils.path_to_croco  # noqa: F401
import croco.utils.misc as misc  # noqa
from croco.utils.misc import NativeScalerWithGradNormCount as NativeScaler  # noqa

from torch.utils.data import default_collate

def get_args_parser():
    parser = argparse.ArgumentParser('DUST3R training', add_help=False)
    # model and criterion
    parser.add_argument('--model', default="AsymmetricCroCo3DStereo(patch_embed_cls='ManyAR_PatchEmbed')",
                        type=str, help="string containing the model to build")
    parser.add_argument('--pretrained', default=None, help='path of a starting checkpoint')
    parser.add_argument('--train_criterion', default="ConfLoss(Regr3D(L21, norm_mode='avg_dis'), alpha=0.2)",
                        type=str, help="train criterion")
    parser.add_argument('--test_criterion', default=None, type=str, help="test criterion")

    # dataset
    parser.add_argument('--train_dataset', required=True, type=str, help="training set")
    parser.add_argument('--test_dataset', default='[None]', type=str, help="testing set")

    # training
    parser.add_argument('--seed', default=0, type=int, help="Random seed")
    parser.add_argument('--batch_size', default=64, type=int,
                        help="Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus")
    parser.add_argument('--accum_iter', default=1, type=int,
                        help="Accumulate gradient iterations (for increasing the effective batch size under memory constraints)")
    parser.add_argument('--epochs', default=800, type=int, help="Maximum number of epochs for the scheduler")

    parser.add_argument('--weight_decay', type=float, default=0.05, help="weight decay (default: 0.05)")
    parser.add_argument('--lr', type=float, default=1e-4, metavar='LR', help='learning rate (absolute lr)')
    parser.add_argument('--blr', type=float, default=1.5e-4, metavar='LR',
                        help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                        help='lower lr bound for cyclic schedulers that hit 0')
    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N', help='epochs to warmup LR')

    parser.add_argument('--amp', type=int, default=0,
                        choices=[0, 1], help="Use Automatic Mixed Precision for pretraining")

    # others
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--allow_first_test', default=1, type=int)
    parser.add_argument('--only_test', default=0, type=int)
    parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes')
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')

    parser.add_argument('--eval_freq', type=int, default=1, help='Test loss evaluation frequency')
    parser.add_argument('--save_freq', default=1, type=int,
                        help='frequence (number of epochs) to save checkpoint in checkpoint-last.pth')
    parser.add_argument('--keep_freq', default=20, type=int,
                        help='frequence (number of epochs) to save checkpoint in checkpoint-%d.pth')
    parser.add_argument('--print_freq', default=20, type=int,
                        help='frequence (number of iterations) to print infos while training')
    parser.add_argument('--miter', default=0, type=int,
                        help='No. of extra inference')
    # output dir
    parser.add_argument('--output_dir', default=None, type=str, help="path where to save the output")
    return parser


def main(args):
    print('args', args)
    misc.init_distributed_mode(args)
    global_rank = misc.get_rank()
    world_size = misc.get_world_size()
    real_batch_size = args.batch_size * world_size
    print('world size', world_size, 'global_rank', global_rank, 'real_batch_size', real_batch_size)
    set_device(args.gpu) # 0

    args.output_dir = get_log_dir_warp(args.output_dir)

    print("output_dir: "+args.output_dir)
    if args.output_dir:
        g_pathmgr.mkdirs(args.output_dir)

    # auto resume
    last_ckpt_fname = os.path.join(args.output_dir, f'checkpoint-last.pth')
    args.resume = last_ckpt_fname if g_pathmgr.isfile(last_ckpt_fname) else None

    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
    print("{}".format(args).replace(', ', ',\n'))

    device = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(device)

    # fix the seed
    seed = args.seed + misc.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)

    cudnn.benchmark = True

    # training dataset and loader
    print('Building train dataset {:s}'.format(args.train_dataset))
    #  dataset and loader
    data_loader_train = build_dataset(args.train_dataset, args.batch_size, args.num_workers, test=False)
    train_epoch_size = real_batch_size * len(data_loader_train)
    print('Building test dataset {:s}'.format(args.test_dataset))
    data_loader_test = {}
    for dataset_name in args.test_dataset.split('+'):
        dataset = build_dataset(dataset_name, args.batch_size, args.num_workers, test=True)
        dataset_name = dataset.dataset.tb_name
        data_loader_test[dataset_name] = dataset

    # model
    print('Loading model: {:s}'.format(args.model))
    model = eval(args.model) # 
    model_name = args.model.split('(')[0]
    print(f'>> Creating train criterion = {args.train_criterion}')
    train_criterion = eval(args.train_criterion).to(device)
    print(f'>> Creating test criterion = {args.test_criterion or args.train_criterion}')
    test_criterion = eval(args.test_criterion or args.criterion).to(device)

    model.to(device)
    model_without_ddp = model
    print("Model = %s" % str(model_without_ddp))

    if args.pretrained and not args.resume:
        model_loaded = eval(model_name).from_pretrained(get_local_path(args.pretrained)).to(device)
        print('Loading pretrained: ', args.pretrained, model_name) # 

        state_dict_loaded = model_loaded.state_dict()
        model.load_state_dict(state_dict_loaded, strict=False)
        model_without_ddp = model

    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
    if args.lr is None:  # only base_lr is specified
        args.lr = args.blr * eff_batch_size / 256
    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
    print("actual lr: %.2e" % args.lr)
    print("accumulate grad iterations: %d" % args.accum_iter)
    print("effective batch size: %d" % eff_batch_size)

    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[args.gpu], find_unused_parameters=True, static_graph=True)
        model_without_ddp = model.module
        total_params = sum(p.numel() for p in model_without_ddp.parameters())
        print(f'Total number of parameters: {total_params}') # ≈1B 

    # following timm: set wd as 0 for bias and norm layers
    param_groups = misc.get_parameter_groups(model_without_ddp, args.weight_decay)
    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
    loss_scaler = NativeScaler()

    def write_log_stats(epoch, train_stats, test_stats):
        if misc.is_main_process():
            if log_writer is not None:
                log_writer.flush()

            log_stats = dict(epoch=epoch, **{f'train_{k}': v for k, v in train_stats.items()})
            for test_name in data_loader_test:
                if test_name not in test_stats:
                    continue
                log_stats.update({test_name+'_'+k: v for k, v in test_stats[test_name].items()})

            with g_pathmgr.open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                f.write(json.dumps(log_stats) + "\n")

    def save_model(epoch, fname, best_so_far):
        misc.save_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer,
                        loss_scaler=loss_scaler, epoch=epoch, fname=fname, best_so_far=best_so_far)

    best_so_far = misc.load_model(args=args, model_without_ddp=model_without_ddp,
                                  optimizer=optimizer, loss_scaler=loss_scaler)
    if best_so_far is None:
        best_so_far = float('inf')
    if global_rank == 0 and args.output_dir is not None:
        log_writer = SummaryWriter(log_dir=args.output_dir)
    else:
        log_writer = None

    print(f"Start training for {args.epochs} epochs from {args.start_epoch}") # 
    start_time = time.time()
    train_stats = test_stats = {}
    for epoch in range(args.start_epoch, args.epochs+1):

        t_save = -time.time()
    
        # Save immediately the last checkpoint
        if epoch > args.start_epoch:
            if args.save_freq and epoch % args.save_freq == 0 or epoch == args.epochs:
                save_model(epoch-1, 'last', best_so_far)
        t_save += time.time()

        # Test on multiple datasets
        new_best = False
        # if False:
        if ((epoch == 0 and args.allow_first_test > 0) or (epoch != 0 and args.eval_freq > 0 and epoch % args.eval_freq == 0)) or epoch == 1:
            test_stats = {}
            test_set_id = -1
            for test_name, testset in data_loader_test.items():
                test_set_id += 1
                t_test = time.time()
                print('test name', test_name)
                stats = test_one_epoch(model, test_criterion, testset,
                                       device, epoch, train_epoch_size, log_writer=log_writer, args=args, prefix=test_name, miter = args.miter, test_set_id = test_set_id)
                test_stats[test_name] = stats

                # Save best of all
                if stats['loss_med'] < best_so_far:
                    best_so_far = stats['loss_med']
                    new_best = True
                print('test epoch time', epoch, time.time() - t_test)

        t_save -= time.time()
        # Save more stuff
        write_log_stats(epoch, train_stats, test_stats)

        if epoch > args.start_epoch:
            if args.keep_freq and epoch % args.keep_freq == 0:
                save_model(epoch-1, str(epoch), best_so_far)
            if new_best:
                save_model(epoch-1, 'best', best_so_far)
        t_save += time.time()

        if epoch >= args.epochs or args.only_test:
            break  # exit after writing last test to disk

        # Train
        t_train = time.time()
        train_stats = train_one_epoch(
            model, train_criterion, data_loader_train,
            optimizer, device, epoch, loss_scaler, train_epoch_size,
            log_writer=log_writer,
            args=args)
        print('train epoch time', epoch, time.time() - t_train)
        print('save epoch time', epoch, t_save)

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))

    save_final_model(args, args.epochs, model_without_ddp, best_so_far=best_so_far)


def save_final_model(args, epoch, model_without_ddp, best_so_far=None):
    checkpoint_path = os.path.join(args.output_dir, 'checkpoint-final.pth')
    to_save = {
        'args': args,
        'model': model_without_ddp if isinstance(model_without_ddp, dict) else model_without_ddp.cpu().state_dict(),
        'epoch': epoch
    }
    if best_so_far is not None:
        to_save['best_so_far'] = best_so_far
    print(f'>> Saving model to {checkpoint_path} ...')
    misc.save_on_master(to_save, checkpoint_path)


def build_dataset(dataset, batch_size, num_workers, test=False):
    split = ['Train', 'Test'][test]
    print(f'Building {split} Data loader for dataset: ', dataset)
    loader = get_data_loader(dataset,
                             batch_size=batch_size,
                             num_workers=num_workers,
                             pin_mem=True,
                             shuffle=not (test),
                             drop_last=not (test))

    print(f"{split} dataset length: ", len(loader))
    return loader

def add_first_best(loss_details, n_ref):

    # import fbvscode
    # fbvscode.set_trace()
    ldk = list(loss_details.keys())
    for k in ldk:
        if k == 'loss':
            continue
        if "_list" in k:
            x_list = np.array(loss_details[k])
            k_base = k.replace('_list', '')
            x_list = x_list.reshape(-1, n_ref)
            x_first = float(x_list[:, 0].mean())
            x_max = float(np.max(x_list, axis = 1).mean())
            x_min = float(np.min(x_list, axis = 1).mean())
            if k_base+'_first' not in ldk:
                loss_details[k_base+'_first'] = x_first
            # if k_base+'_best' not in ldk:
            loss_details[k_base+'_max'] = x_max
            loss_details[k_base+'_min'] = x_min

    return loss_details

def postprocess_batch(batch): # here the randomized number of inference views / number of rendered views are applied to the whole batch.

    nv, nr = batch[0]['random_nv_nr'][0].cpu().numpy() # we are always using the first sample's No. of views / No. of rendered views and apply it to all samples in the batch
    while len(batch) > nv:
        del batch[-1]
    batch = batch[:nv]
    ni = nv - nr
    for i in range(ni):
        batch[i]['only_render'][:] = False
    for i in range(ni, nv):
        batch[i]['only_render'][:] = True
    return batch, ni

def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
                    data_loader: Sized, optimizer: torch.optim.Optimizer,
                    device: torch.device, epoch: int, loss_scaler, epoch_size,
                    args,
                    log_writer=None):
    t_all = -time.time()
    assert torch.backends.cuda.matmul.allow_tf32 == True
    t_misc_1 = -time.time()
    model.train(True)
    metric_logger = misc.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    header = 'Epoch: [{}]'.format(epoch)
    accum_iter = args.accum_iter

    t_misc_1 += time.time()

    t_misc_2 = -time.time()

    if hasattr(data_loader, 'dataset') and hasattr(data_loader.dataset, 'set_epoch'):
        data_loader.dataset.set_epoch(epoch)
    if hasattr(data_loader, 'sampler') and hasattr(data_loader.sampler, 'set_epoch'):
        data_loader.sampler.set_epoch(epoch)

    optimizer.zero_grad()
    t_misc_2 += time.time()
    t_misc_3 = 0
    t_misc_4 = 0

    t_inference = 0
    t_bp = 0

    print('before training')

    t_all_time = [time.time()]
    for data_iter_step, batch in enumerate(metric_logger.log_every(data_loader, args.print_freq, header)):
        batch, ni = postprocess_batch(batch)
        t_misc_3 -= time.time()
        epoch_f = epoch + data_iter_step / len(data_loader)

        # we use a per iteration (instead of per epoch) lr scheduler
        if data_iter_step % accum_iter == 0:
            misc.adjust_learning_rate(optimizer, epoch_f, args)
        
        t_misc_3 += time.time()
        
        t_inference_i = -time.time()
        print('check sync train before foward', misc.get_rank(), epoch, data_iter_step)
        # torch.cuda.synchronize()
        # delta = 1
        # if epoch > 0:
        #     delta = 0
        delta = 1
        need_log = data_iter_step == 0 or ((data_iter_step + delta) % accum_iter == 0 and ((data_iter_step + delta) % (accum_iter * args.print_freq)) == 0)

        loss_tuple = loss_of_one_batch(batch, model, criterion, device,
                                       symmetrize_batch=True,
                                       use_amp=bool(args.amp), ret='loss', log = need_log)
        # torch.cuda.synchronize()
        t_inference_i += time.time()
        t_inference += t_inference_i
        t_bp_i = -time.time()
        loss, loss_details = loss_tuple  # criterion returns two values

        print('check sync train after forward', misc.get_rank(), epoch, data_iter_step)
        loss /= accum_iter
        if loss > 10:
            print('strange loss appears', loss)
            loss = loss * 0.
        norm = loss_scaler(loss, optimizer, parameters=model.parameters(), # backward inside, no clip grad
                    update_grad=(data_iter_step + 1) % accum_iter == 0, model = model)
        if norm is not None and norm > 1000:
            print('strange norm appears', norm)
        if (data_iter_step + 1) % accum_iter == 0:
            optimizer.zero_grad()
        t_bp_i += time.time()
        t_bp += t_bp_i
        t_all_time.append(time.time())
        print('train batch', 'step', data_iter_step, 'len data', len(data_loader), 'rank', misc.get_rank(), 'epoch_f', epoch_f, 'inference time', t_inference_i, t_inference / (1 + data_iter_step),  'bp time', t_bp_i, t_bp / (1 + data_iter_step), 'all time', t_all_time[-1] - t_all_time[-2], (t_all_time[-1] - t_all_time[0]) / (1 + data_iter_step))
        # inference time 0.23065853118896484 bp time 0.42483043670654297 # all time is similar for 4x8 and 1x8, which means 4x8 is indeed more efficient
        t_misc_4 -= time.time()
        lr = optimizer.param_groups[0]["lr"]
        for k in list(loss_details.keys()):
            if not isinstance(loss_details[k], (float, int)):
                loss_details.pop(k)
        if need_log:
            
            loss_value = float(loss * accum_iter)
            if not math.isfinite(loss_value):
                print("Loss is {}, stopping training".format(loss_value), force=True)
                sys.exit(1)
            metric_logger.update(epoch=epoch_f)
            metric_logger.update(lr=lr)
            metric_logger.update(loss=loss_value, **loss_details)

        del loss
        del batch
        
        # print('train_loss debug', data_iter_step, accum_iter, data_iter_step, args.print_freq, ((data_iter_step + 1) % (accum_iter * args.print_freq)), log_writer)
        if need_log:
            loss_value_reduce = misc.all_reduce_mean(loss_value)  # MUST BE EXECUTED BY ALL NODES
            if log_writer is None:
                continue
            """ We use epoch_1000x as the x-axis in tensorboard.
            This calibrates different curves when batch size changes.
            """
            # epoch_1000x = int(epoch_f * 1000)
            epoch_1000x = int(epoch_f * epoch_size)
            log_writer.add_scalar('train_loss', loss_value_reduce, epoch_1000x)
            norm_item = norm.item() if norm is not None else 0.
            log_writer.add_scalar('train_grad_norm', norm_item, epoch_1000x)
            log_writer.add_scalar('train/time/all', t_all_time[-1] - t_all_time[-2], epoch_1000x)
            log_writer.add_scalar('train/time/ff', t_inference_i, epoch_1000x)
            log_writer.add_scalar('train/time/bp', t_bp_i, epoch_1000x)
            log_writer.add_scalar('train_lr', lr, epoch_1000x)
            log_writer.add_scalar('train_iter', epoch_f * len(data_loader), epoch_1000x)
            for name, val in loss_details.items():
                log_writer.add_scalar('train_'+name, val, epoch_1000x)
        t_misc_4 += time.time()
    # gather the stats from all processes
    t_misc_5 = -time.time()
    metric_logger.synchronize_between_processes()
    t_misc_5 += time.time()

    t_all += time.time()
    print('train misc time', t_misc_1, t_misc_2, t_misc_3, t_misc_4, t_misc_5, t_all, t_inference, t_bp, t_all - t_inference - t_bp) # all miscs are very small train misc time 0.041296958923339844 0.0005085468292236328 0.002261638641357422 0.0012340545654296875 130.9805166721344
    print("Averaged stats:", metric_logger)
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

def save_results(loss_and_others, batch, name_list, args):
    all_info = loss_and_others
    other_info = loss_and_others['loss'][1]
    # view1: img (real_bs * 2 (data aug for symmetry), 3, res=224, res), depthmap, camera_pose (real_bs * 2, 4, 4), camera_intrinsics, dataset, label, instance, idx, true_shape, pts3d (real_bs * 2, res, res, 3), valid_mask, rng
    # pred1: pts3d, conf
    # pred2: pts3d_in_other_view, conf
    g_pathmgr.mkdirs(args.output_dir + '/results')
    g_pathmgr.mkdirs(args.output_dir + '/videos')
    bs = all_info['view1']['img'].shape[0] # real_bs * 2 = bs
    
    if 'view2s' in all_info.keys(): # MV here
        n_ref = all_info['pred1']['pts3d'].shape[0] // bs
        if n_ref != 1:
            from dust3r.losses import extend_gts
            views = [all_info['view1']] + all_info['view2s']
            views = extend_gts(views, n_ref, bs)
            all_info['view1'] = views[0]
            all_info['view2s'] = views[1:]
            bs = n_ref * bs

        for img_id in range(bs):
            # import fbvscode
            # fbvscode.set_trace()
            img_id_mref_first = img_id
            # img_id_mref_first = n_ref * img_id # 00022_id_000000001_test_dataName_hs_3.0_sceneName_Beach_refId_00_00000_0033_test
            label = batch[0]['label'][img_id // n_ref]
            name = "_".join(name_list[0:1] + [label] + name_list[1:])
            rgb1 = all_info['view1']['img'][img_id].permute(1,2,0)
            valid_mask1 = all_info['view1']['valid_mask'][img_id].reshape(-1)
            num_render_views = all_info['view2s'][0].get("num_render_views", torch.zeros([0]).long())[0].item()
            rgb2s_all = [x['img'][img_id].permute(1,2,0) for x in all_info['view2s']]
            valid_mask2s = [x['valid_mask'][img_id].reshape(-1) for x in all_info['view2s']]
            rgb2s = rgb2s_all[:-num_render_views] if num_render_views else rgb2s_all
            valid_mask2s = valid_mask2s[:-num_render_views] if num_render_views else valid_mask2s
            rgb      = torch.cat([rgb1.reshape(-1, 3)] + [rgb2.reshape(-1, 3) for rgb2 in rgb2s], 0)
            valid_masks = torch.stack([valid_mask1] + valid_mask2s, 0)
            pts3d_gt = torch.cat([all_info['view1']['pts3d'][img_id].reshape(-1, 3)] + [x['pts3d'][img_id].reshape(-1, 3) for x in (all_info['view2s'][:-num_render_views] if num_render_views else all_info['view2s'])], 0)
            pts3d    = torch.cat([all_info['pred1']['pts3d'][img_id_mref_first].reshape(-1, 3)] + [x['pts3d_in_other_view'][img_id_mref_first].reshape(-1, 3) for x in all_info['pred2s']], 0)
            conf = torch.cat([all_info['pred1']['conf'][img_id_mref_first].reshape(-1, 1)] + [x['conf'][img_id_mref_first].reshape(-1, 1) for x in all_info['pred2s']], 0) # [N, 1]
            conf_sorted = conf.reshape(-1).sort()[0]
            conf_thres = float(conf_sorted[int(conf.shape[0] * 0.03)])
            cam1 = all_info['view1']['camera_pose'][img_id] # c2w
            pts3d = geotrf(cam1, pts3d)  # B,H,W,3
            # img_id_name = str(img_id).zfill(3)
            img_id_name = str(time.time()).split('.')[1]
            img_id_name = f"nref_{img_id % n_ref}_{str(time.time()).split('.')[1]}"
            video_pcd_gt =      pcd_render(pts3d_gt, rgb, tgt = None, normalize = True)
            video_pcd =         pcd_render(pts3d   , rgb, tgt = None, normalize = True)
            video_pcd_conf =    pcd_render(pts3d   , rgb, tgt = None, normalize = True, mask = (conf > conf_thres) * valid_masks.reshape(-1, 1)) # log(3)
            print('vis conf range', conf.min(), conf.mean(), conf.max(), conf_thres, (conf < 1.02).float().mean(), (conf < 1.03).float().mean(), (conf < 1.06).float().mean(), (conf < 1.09).float().mean())
            save_video_combined([video_pcd, video_pcd_conf, video_pcd_gt], f"{args.output_dir}/videos/{name}_{img_id_name}_and_gt.mp4")
            if 'scale' in all_info['pred1'].keys(): # 3DGS predicted
                gts = [all_info['view1']] + [v for v in (all_info['view2s'][:-num_render_views] if num_render_views else all_info['view2s'])]
                preds = [all_info['pred1']] + [v for v in all_info['pred2s']]
                video_gs_gt = gs_render(gts, preds, img_id, img_id_mref_first, cam1, normalize = True, gt_pcd = True, gt_img = True)
                video_gs_gt_img_only = gs_render(gts, preds, img_id, img_id_mref_first, cam1, normalize = True, gt_pcd = False, gt_img = True)
                video_gs =    gs_render(gts, preds, img_id, img_id_mref_first, cam1, normalize = True)
                save_video_combined([video_gs, video_gs_gt_img_only, video_gs_gt], f"{args.output_dir}/videos/{name}_{img_id_name}_and_gt_GS.mp4")
            # import fbvscode
            # fbvscode.set_trace()
            other_info_web = {k: float(other_info[k][img_id_mref_first]) for k in other_info.keys() if "_list" in k}
            torch.save(other_info_web, f"{args.output_dir}/videos/{name}_{img_id_name}.pth")
            # rgb is -1~1, shape = (res,res,3)
            rgbs = [rgb1]
            save_image_manifold(((rgb1 + 1) / 2 * 255).cpu().numpy().astype(np.uint8), f"{args.output_dir}/videos/{name}_{img_id_name}_rgb1.png")
            for rgb_id, rgb2 in enumerate(rgb2s_all):
                rgbs.append(rgb2)
                save_image_manifold(((rgb2 + 1) / 2 * 255).cpu().numpy().astype(np.uint8), f"{args.output_dir}/videos/{name}_{img_id_name}_rgb{rgb_id + 2}.png")
            rgbs = torch.cat(rgbs, dim = 1) # [h,w (combine here),3]
            save_image_manifold(((rgbs + 1) / 2 * 255).cpu().numpy().astype(np.uint8), f"{args.output_dir}/videos/{name}_{img_id_name}_rgb_all.png")
            if "render_all" in other_info.keys():
                render_all = other_info["render_all"] # render_all[img_id]: [nv, 224, 224, 3]
                save_image_manifold(((render_all[img_id_mref_first].permute(1,0,2,3).flatten(1,2) + 1) / 2 * 255).cpu().numpy().astype(np.uint8), f"{args.output_dir}/videos/{name}_{img_id_name}_gs.png")
            if "render_relocated_all" in other_info.keys():
                render_relocated_all = other_info["render_relocated_all"] # render_all[img_id]: [nv, 224, 224, 3]
                save_image_manifold(((render_relocated_all[img_id_mref_first].permute(1,0,2,3).flatten(1,2) + 1) / 2 * 255).cpu().numpy().astype(np.uint8), f"{args.output_dir}/videos/{name}_{img_id_name}_gs_relocated.png")
    else:
        for img_id in range(bs):
            rgb1 = all_info['view1']['img'][img_id].permute(1,2,0)
            rgb2 = all_info['view2']['img'][img_id].permute(1,2,0)
            rgb      = torch.cat([rgb1.reshape(-1, 3), rgb2.reshape(-1, 3)], 0)
            pts3d_gt = torch.cat([all_info['view1']['pts3d'][img_id].reshape(-1, 3), all_info['view2']['pts3d'][img_id].reshape(-1, 3)], 0)
            pts3d    = torch.cat([all_info['pred1']['pts3d'][img_id].reshape(-1, 3), all_info['pred2']['pts3d_in_other_view'][img_id].reshape(-1, 3)], 0)

            cam1 = all_info['view1']['camera_pose'][img_id] # c2w -> w2c
            pts3d = geotrf(cam1, pts3d)  # B,H,W,3
            
            img_id_name = str(img_id).zfill(3)
            pcd_render(pts3d_gt, rgb, f"{args.output_dir}/videos/{name}_{img_id_name}_gt.mp4", normalize = True)
            pcd_render(pts3d   , rgb, f"{args.output_dir}/videos/{name}_{img_id_name}.mp4", normalize = True)
            torch.save(loss_and_others, f"{args.output_dir}/results/{name}_{img_id_name}.pth")
            # rgb is -1~1, shape = (res,res,3)
            save_image_manifold(((rgb1 + 1) / 2 * 255).cpu().numpy().astype(np.uint8), f"{args.output_dir}/videos/{name}_{img_id_name}_rgb1.png")
            save_image_manifold(((rgb2 + 1) / 2 * 255).cpu().numpy().astype(np.uint8), f"{args.output_dir}/videos/{name}_{img_id_name}_rgb2.png")

def to_device(data, device):

    if torch.is_tensor(data):
        return data.to(device)
    elif isinstance(data, dict):
        return {key: to_device(value, device) for key, value in data.items()}
    elif isinstance(data, list):
        return [to_device(element, device) for element in data]
    elif isinstance(data, tuple):
        return tuple(to_device(element, device) for element in data)
    else:
        return data

def update_batch(batch, loss_and_others, data_loader):
    views = [loss_and_others['view1']] + loss_and_others['view2s']
    ids = views[0]['idx'][0]
    pts_pred = [loss_and_others['pred1']['pts3d']] + [x['pts3d_in_other_view'] for x in loss_and_others['pred2s']] # [bs, res, res, 3] each
    pts_pred = torch.stack(pts_pred, dim = 1) # [bs, n_inference, res, res, 3]
    pts_pred_center_view = pts_pred.mean(dim = (2,3)) # [bs, n_inference, 3]
    pts_pred_center = pts_pred_center_view.mean(dim = 1) # [bs, 3]
    view_dis = torch.norm(pts_pred_center_view - pts_pred_center.unsqueeze(1), dim = 2) # [bs, n_inference]
    nearest_view_id = view_dis.argmin(dim = 1) # [bs]
    new_batch = [data_loader.dataset.__getitem_bsvd__(x.item(), y.item()) for x, y in zip(ids.long(), nearest_view_id)]
    new_batch = default_collate(new_batch)
    return new_batch

@torch.no_grad()
def test_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
                   data_loader: Sized, device: torch.device, epoch: int,
                   train_epoch_size, args, log_writer=None, prefix='test', miter = False, test_set_id = 0):
    t_begin1 = -time.time()
    model.eval()
    metric_logger = misc.MetricLogger(delimiter="  ")
    metric_logger.meters = defaultdict(lambda: misc.SmoothedValue(window_size=9**9))
    header = 'Test Epoch: [{}]'.format(epoch)

    if log_writer is not None:
        print('log_dir: {}'.format(log_writer.log_dir))
    t_begin1 += time.time()
    t_begin2 = -time.time()
    if hasattr(data_loader, 'dataset') and hasattr(data_loader.dataset, 'set_epoch'):
        print('set in dataset')
        data_loader.dataset.set_epoch(epoch)
    if hasattr(data_loader, 'sampler') and hasattr(data_loader.sampler, 'set_epoch'):
        print('set in sampler')
        data_loader.sampler.set_epoch(epoch)
    t_begin2 += time.time()
    t_batch = -time.time()
    t_inference = 0
    t_save = 0
    for batch_id, batch in enumerate(metric_logger.log_every(data_loader, args.print_freq, header)):
        batch, ni = postprocess_batch(batch)
        t = time.time()
        # print('test batch 1st', batch_id, len(data_loader), epoch, misc.get_rank())
        # torch.cuda.synchronize()
        t_inference -= time.time()
        if miter:
            loss_and_others = loss_of_one_batch(batch, model, criterion, device,
                                       symmetrize_batch=True,
                                       use_amp=bool(args.amp), ret=None)
            batch = update_batch(batch, loss_and_others, data_loader)
            # print('pts3d_2_avg before', loss_and_others['loss'][1]['Regr3D_ScaleShiftInv_pts3d_2'])
            # import fbvscode
            # fbvscode.set_trace()
        loss_and_others = loss_of_one_batch(batch, model, criterion, device,
                                       symmetrize_batch=True,
                                       use_amp=bool(args.amp), ret=None)
        # print('pts3d_2_avg after', loss_and_others['loss'][1]['Regr3D_ScaleShiftInv_pts3d_2'])
        # torch.cuda.synchronize()
        t_inference += time.time()
        # print('test batch 2nd', batch_id, len(data_loader), epoch, misc.get_rank(), 'inference time', time.time() - t, batch[0]['pts3d'].shape)

        t_save -= time.time()
        # print('data_loader', type(data_loader.dataset).__name__)
        if data_loader.dataset.save_results:
            global_rank = misc.get_rank()
            prefix_save = [str(epoch).zfill(5) + "_testSetID_" + str(test_set_id).zfill(3)]
            # prefix_save = [str(epoch).zfill(5), str(batch_id).zfill(5), str(global_rank).zfill(4), data_loader.dataset.save_prefix]
            save_results(loss_and_others, batch, prefix_save, args)
        t_save += time.time()
        # print('test batch 3rd', batch_id, len(data_loader), epoch, misc.get_rank())
        loss_tuple = loss_and_others['loss']
        loss_value, loss_details = loss_tuple  # criterion returns two values
        n_ref = int(loss_details['n_ref'])
        loss_details.pop('n_ref')
        loss_details = add_first_best(loss_details, n_ref)
        for k in list(loss_details.keys()):
            if not isinstance(loss_details[k], (float, int)):
                loss_details.pop(k)
        metric_logger.update(loss=float(loss_value), **loss_details)
        # if batch_id >= 1:
        #     break

    t_batch += time.time()
    # gather the stats from all processes
    t_log = - time.time()
    
    if data_loader.dataset.save_results and misc.get_rank() == 0:
        if generate_html is not None:
            generate_html(args.output_dir + '/videos', args.output_dir + '/html')

    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)

    aggs = [('avg', 'global_avg'), ('med', 'median')]
    results = {f'{k}_{tag}': getattr(meter, attr) for k, meter in metric_logger.meters.items() for tag, attr in aggs}

    if log_writer is not None:
        for name, val in results.items():
            # epoch_1000x = int(epoch * 1000)
            epoch_1000x = int(epoch * train_epoch_size)
            log_writer.add_scalar(prefix+'/'+name, val, epoch_1000x)
    t_log += time.time()
    print('test all time', prefix, 'batch', t_batch, t_batch - t_inference - t_save, 'inference', t_inference, 'save', t_save, 'log', t_log, 'two begins', t_begin1, t_begin2) # inference and log is small, batch is kind of large, but  
    # test all time  100 @ ScannetPair_test batch 70.40310192108154 inference 5.6025426387786865 save 0.0006468296051025391 log 0.0017290115356445312
    # seems batch cost a lot of time, maybe from dataloading? testing now, inference is fast, save cost time in visualization but not torch.save, t_log and t_begin is fast.
    return results


if __name__ == '__main__':
    args = get_args_parser()
    args = args.parse_args()
    main(args)