vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/duster/dust3r/training.py

@@ -0,0 +1,377 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# training code for DUSt3R
+# --------------------------------------------------------
+# References:
+# MAE: https://github.com/facebookresearch/mae
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import argparse
+import datetime
+import json
+import numpy as np
+import os
+import sys
+import time
+import math
+from collections import defaultdict
+from pathlib import Path
+from typing import Sized
+
+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, 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
+
+import dust3r.utils.path_to_croco  # noqa: F401
+import croco.utils.misc as misc  # noqa
+from croco.utils.misc import NativeScalerWithGradNormCount as NativeScaler  # noqa
+
+
+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=None, 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")
+    parser.add_argument("--disable_cudnn_benchmark", action='store_true', default=False,
+                        help="set cudnn.benchmark = False")
+    # others
+    parser.add_argument('--num_workers', default=8, 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')
+
+    # output dir
+    parser.add_argument('--output_dir', default='./output/', type=str, help="path where to save the output")
+    return parser
+
+
+def train(args):
+    misc.init_distributed_mode(args)
+    global_rank = misc.get_rank()
+    world_size = misc.get_world_size()
+
+    print("output_dir: " + args.output_dir)
+    if args.output_dir:
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+    # auto resume
+    last_ckpt_fname = os.path.join(args.output_dir, f'checkpoint-last.pth')
+    args.resume = last_ckpt_fname if os.path.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 = not args.disable_cudnn_benchmark
+
+    # 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)
+    print('Building test dataset {:s}'.format(args.train_dataset))
+    data_loader_test = {dataset.split('(')[0]: build_dataset(dataset, args.batch_size, args.num_workers, test=True)
+                        for dataset in args.test_dataset.split('+')}
+
+    # model
+    print('Loading model: {:s}'.format(args.model))
+    model = eval(args.model)
+    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:
+        print('Loading pretrained: ', args.pretrained)
+        ckpt = torch.load(args.pretrained, map_location=device)
+        print(model.load_state_dict(ckpt['model'], strict=False))
+        del ckpt  # in case it occupies memory
+
+    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
+
+    # 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))
+    print(optimizer)
+    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 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")
+    start_time = time.time()
+    train_stats = test_stats = {}
+    for epoch in range(args.start_epoch, args.epochs + 1):
+
+        # 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)
+
+        # Test on multiple datasets
+        new_best = False
+        if (epoch > 0 and args.eval_freq > 0 and epoch % args.eval_freq == 0):
+            test_stats = {}
+            for test_name, testset in data_loader_test.items():
+                stats = test_one_epoch(model, test_criterion, testset,
+                                       device, epoch, log_writer=log_writer, args=args, prefix=test_name)
+                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
+
+        # 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)
+        if epoch >= args.epochs:
+            break  # exit after writing last test to disk
+
+        # Train
+        train_stats = train_one_epoch(
+            model, train_criterion, data_loader_train,
+            optimizer, device, epoch, loss_scaler,
+            log_writer=log_writer,
+            args=args)
+
+    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):
+    output_dir = Path(args.output_dir)
+    checkpoint_path = 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 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,
+                    args,
+                    log_writer=None):
+    assert torch.backends.cuda.matmul.allow_tf32 == True
+
+    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
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+
+    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()
+
+    for data_iter_step, batch in enumerate(metric_logger.log_every(data_loader, args.print_freq, header)):
+        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)
+
+        loss_tuple = loss_of_one_batch(batch, model, criterion, device,
+                                       symmetrize_batch=True,
+                                       use_amp=bool(args.amp), ret='loss')
+        loss, loss_details = loss_tuple  # criterion returns two values
+        loss_value = float(loss)
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value), force=True)
+            sys.exit(1)
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(),
+                    update_grad=(data_iter_step + 1) % accum_iter == 0)
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        del loss
+        del batch
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(epoch=epoch_f)
+        metric_logger.update(lr=lr)
+        metric_logger.update(loss=loss_value, **loss_details)
+
+        if (data_iter_step + 1) % accum_iter == 0 and ((data_iter_step + 1) % (accum_iter * args.print_freq)) == 0:
+            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)
+            log_writer.add_scalar('train_loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('train_lr', lr, epoch_1000x)
+            log_writer.add_scalar('train_iter', epoch_1000x, epoch_1000x)
+            for name, val in loss_details.items():
+                log_writer.add_scalar('train_' + name, val, epoch_1000x)
+
+    # 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 test_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
+                   data_loader: Sized, device: torch.device, epoch: int,
+                   args, log_writer=None, prefix='test'):
+
+    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))
+
+    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)
+
+    for _, batch in enumerate(metric_logger.log_every(data_loader, args.print_freq, header)):
+        loss_tuple = loss_of_one_batch(batch, model, criterion, device,
+                                       symmetrize_batch=True,
+                                       use_amp=bool(args.amp), ret='loss')
+        loss_value, loss_details = loss_tuple  # criterion returns two values
+        metric_logger.update(loss=float(loss_value), **loss_details)
+
+    # gather the stats from all processes
+    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():
+            log_writer.add_scalar(prefix + '_' + name, val, 1000 * epoch)
+
+    return results

vismatch/third_party/duster/dust3r/utils/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).

vismatch/third_party/duster/dust3r/utils/device.py

@@ -0,0 +1,76 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# utilitary functions for DUSt3R
+# --------------------------------------------------------
+import numpy as np
+import torch
+
+
+def todevice(batch, device, callback=None, non_blocking=False):
+    ''' Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy).
+
+    batch: list, tuple, dict of tensors or other things
+    device: pytorch device or 'numpy'
+    callback: function that would be called on every sub-elements.
+    '''
+    if callback:
+        batch = callback(batch)
+
+    if isinstance(batch, dict):
+        return {k: todevice(v, device) for k, v in batch.items()}
+
+    if isinstance(batch, (tuple, list)):
+        return type(batch)(todevice(x, device) for x in batch)
+
+    x = batch
+    if device == 'numpy':
+        if isinstance(x, torch.Tensor):
+            x = x.detach().cpu().numpy()
+    elif x is not None:
+        if isinstance(x, np.ndarray):
+            x = torch.from_numpy(x)
+        if torch.is_tensor(x):
+            x = x.to(device, non_blocking=non_blocking)
+    return x
+
+
+to_device = todevice  # alias
+
+
+def to_numpy(x): return todevice(x, 'numpy')
+def to_cpu(x): return todevice(x, 'cpu')
+def to_cuda(x): return todevice(x, 'cuda')
+
+
+def collate_with_cat(whatever, lists=False):
+    if isinstance(whatever, dict):
+        return {k: collate_with_cat(vals, lists=lists) for k, vals in whatever.items()}
+
+    elif isinstance(whatever, (tuple, list)):
+        if len(whatever) == 0:
+            return whatever
+        elem = whatever[0]
+        T = type(whatever)
+
+        if elem is None:
+            return None
+        if isinstance(elem, (bool, float, int, str)):
+            return whatever
+        if isinstance(elem, tuple):
+            return T(collate_with_cat(x, lists=lists) for x in zip(*whatever))
+        if isinstance(elem, dict):
+            return {k: collate_with_cat([e[k] for e in whatever], lists=lists) for k in elem}
+
+        if isinstance(elem, torch.Tensor):
+            return listify(whatever) if lists else torch.cat(whatever)
+        if isinstance(elem, np.ndarray):
+            return listify(whatever) if lists else torch.cat([torch.from_numpy(x) for x in whatever])
+
+        # otherwise, we just chain lists
+        return sum(whatever, T())
+
+
+def listify(elems):
+    return [x for e in elems for x in e]