import os from pathlib import Path import sys if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: sys.path.insert(0, _package_root) import json import time import random from typing import * import itertools from contextlib import nullcontext from concurrent.futures import ThreadPoolExecutor import io import numpy as np import cv2 from PIL import Image import torch import torch.nn as nn import torch.nn.functional as F import torch.version import accelerate from accelerate import Accelerator, DistributedDataParallelKwargs from accelerate.utils import set_seed import utils3d import click from tqdm import tqdm, trange import mlflow torch.backends.cudnn.benchmark = False # Varying input size, make sure cudnn benchmark is disabled from moge.train.dataloader import TrainDataLoaderPipeline from moge.train.losses import ( affine_invariant_global_loss, affine_invariant_local_loss, edge_loss, normal_loss, mask_l2_loss, mask_bce_loss, metric_scale_loss, normal_map_loss, monitoring, ) from moge.train.utils import build_optimizer, build_lr_scheduler from moge.utils.geometry_torch import intrinsics_to_fov from moge.utils.vis import colorize_depth, colorize_normal from moge.utils.tools import key_average, recursive_replace, CallbackOnException, flatten_nested_dict from moge.test.metrics import compute_metrics @click.command() @click.option('--config', 'config_path', type=str, default='configs/debug.json') @click.option('--workspace', type=str, default='workspace/debug', help='Path to the workspace') @click.option('--checkpoint', 'checkpoint_path', type=str, default=None, help='Path to the checkpoint to load. "latest" to load latest checkpoint in workspace, integer to load by step number') @click.option('--batch_size_forward', type=int, default=8, help='Batch size for each forward pass on each device') @click.option('--gradient_accumulation_steps', type=int, default=1, help='Number of steps to accumulate gradients') @click.option('--enable_gradient_checkpointing', type=bool, default=True, help='Use gradient checkpointing in backbone') @click.option('--enable_mixed_precision', type=bool, default=False, help='Use mixed precision training. Backbone is converted to FP16') @click.option('--enable_ema', type=bool, default=True, help='Maintain an exponential moving average of the model weights') @click.option('--num_iterations', type=int, default=1000000, help='Number of iterations to train the model') @click.option('--save_every', type=int, default=10000, help='Save checkpoint every n iterations') @click.option('--log_every', type=int, default=1000, help='Log metrics every n iterations') @click.option('--vis_every', type=int, default=0, help='Visualize every n iterations') @click.option('--num_vis_images', type=int, default=32, help='Number of images to visualize, must be a multiple of divided batch size') @click.option('--enable_mlflow', type=bool, default=True, help='Log metrics to MLFlow') @click.option('--seed', type=int, default=0, help='Random seed') def main( config_path: str, workspace: str, checkpoint_path: str, batch_size_forward: int, gradient_accumulation_steps: int, enable_gradient_checkpointing: bool, enable_mixed_precision: bool, enable_ema: bool, num_iterations: int, save_every: int, log_every: int, vis_every: int, num_vis_images: int, enable_mlflow: bool, seed: Optional[int], ): # Load config with open(config_path, 'r') as f: config = json.load(f) accelerator = Accelerator( gradient_accumulation_steps=gradient_accumulation_steps, mixed_precision='fp16' if enable_mixed_precision else None, kwargs_handlers=[ DistributedDataParallelKwargs(find_unused_parameters=True) ] ) device = accelerator.device batch_size_total = batch_size_forward * gradient_accumulation_steps * accelerator.num_processes # Log config if accelerator.is_main_process: if enable_mlflow: try: mlflow.log_params({ **click.get_current_context().params, 'batch_size_total': batch_size_total, }) except: print('Failed to log config to MLFlow') Path(workspace).mkdir(parents=True, exist_ok=True) with Path(workspace).joinpath('config.json').open('w') as f: json.dump(config, f, indent=4) # Set seed if seed is not None: set_seed(seed, device_specific=True) # Initialize model print('Initialize model') with accelerator.local_main_process_first(): from moge.model import import_model_class_by_version MoGeModel = import_model_class_by_version(config['model_version']) model = MoGeModel(**config['model']) count_total_parameters = sum(p.numel() for p in model.parameters()) print(f'Total parameters: {count_total_parameters}') # Set up EMA model if enable_ema and accelerator.is_main_process: ema_avg_fn = lambda averaged_model_parameter, model_parameter, num_averaged: 0.999 * averaged_model_parameter + 0.001 * model_parameter ema_model = torch.optim.swa_utils.AveragedModel(model, device=accelerator.device, avg_fn=ema_avg_fn) # Set gradient checkpointing if enable_gradient_checkpointing: model.enable_gradient_checkpointing() import warnings warnings.filterwarnings("ignore", category=FutureWarning, module="torch.utils.checkpoint") # Initalize optimizer & lr scheduler optimizer = build_optimizer(model, config['optimizer']) lr_scheduler = build_lr_scheduler(optimizer, config['lr_scheduler']) count_grouped_parameters = [sum(p.numel() for p in param_group['params'] if p.requires_grad) for param_group in optimizer.param_groups] for i, count in enumerate(count_grouped_parameters): print(f'- Group {i}: {count} parameters') # Attempt to load checkpoint checkpoint: Dict[str, Any] with accelerator.local_main_process_first(): if checkpoint_path is None: # - No checkpoint checkpoint = None elif checkpoint_path.endswith('.pt'): # - Load specific checkpoint file print(f'Load checkpoint: {checkpoint_path}') checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True) elif checkpoint_path == "latest": # - Load latest checkpoint checkpoint_path = Path(workspace, 'checkpoint', 'latest.pt') if checkpoint_path.exists(): print(f'Load checkpoint: {checkpoint_path}') checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True) i_step = checkpoint['step'] if 'model' not in checkpoint and (checkpoint_model_path := Path(workspace, 'checkpoint', f'{i_step:08d}.pt')).exists(): print(f'Load model checkpoint: {checkpoint_model_path}') checkpoint['model'] = torch.load(checkpoint_model_path, map_location='cpu', weights_only=True)['model'] if 'optimizer' not in checkpoint and (checkpoint_optimizer_path := Path(workspace, 'checkpoint', f'{i_step:08d}_optimizer.pt')).exists(): print(f'Load optimizer checkpoint: {checkpoint_optimizer_path}') checkpoint.update(torch.load(checkpoint_optimizer_path, map_location='cpu', weights_only=True)) if enable_ema and accelerator.is_main_process: if 'ema_model' not in checkpoint and (checkpoint_ema_model_path := Path(workspace, 'checkpoint', f'{i_step:08d}_ema.pt')).exists(): print(f'Load EMA model checkpoint: {checkpoint_ema_model_path}') checkpoint['ema_model'] = torch.load(checkpoint_ema_model_path, map_location='cpu', weights_only=True)['model'] else: print(f'No latest checkpoint found. Start from scratch.') checkpoint = None else: # - Load by step number i_step = int(checkpoint_path) checkpoint = {'step': i_step} if (checkpoint_model_path := Path(workspace, 'checkpoint', f'{i_step:08d}.pt')).exists(): print(f'Load model checkpoint: {checkpoint_model_path}') checkpoint['model'] = torch.load(checkpoint_model_path, map_location='cpu', weights_only=True)['model'] if (checkpoint_optimizer_path := Path(workspace, 'checkpoint', f'{i_step:08d}_optimizer.pt')).exists(): print(f'Load optimizer checkpoint: {checkpoint_optimizer_path}') checkpoint.update(torch.load(checkpoint_optimizer_path, map_location='cpu', weights_only=True)) if enable_ema and accelerator.is_main_process: if (checkpoint_ema_model_path := Path(workspace, 'checkpoint', f'{i_step:08d}_ema.pt')).exists(): print(f'Load EMA model checkpoint: {checkpoint_ema_model_path}') checkpoint['ema_model'] = torch.load(checkpoint_ema_model_path, map_location='cpu', weights_only=True)['model'] if checkpoint is None: # Initialize model weights print('Initialize model weights') with accelerator.local_main_process_first(): model.init_weights() initial_step = 0 else: model.load_state_dict(checkpoint['model'], strict=False) if 'step' in checkpoint: initial_step = checkpoint['step'] + 1 else: initial_step = 0 if 'optimizer' in checkpoint: optimizer.load_state_dict(checkpoint['optimizer']) if enable_ema and accelerator.is_main_process and 'ema_model' in checkpoint: ema_model.module.load_state_dict(checkpoint['ema_model'], strict=False) if 'lr_scheduler' in checkpoint: lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) del checkpoint model, optimizer = accelerator.prepare(model, optimizer) if torch.version.hip and isinstance(model, torch.nn.parallel.DistributedDataParallel): # Hacking potential gradient synchronization issue in ROCm backend from moge.model.utils import sync_ddp_hook model.register_comm_hook(None, sync_ddp_hook) # Initialize training data pipeline with accelerator.local_main_process_first(): train_data_pipe = TrainDataLoaderPipeline(config['data'], batch_size_forward) def _write_bytes_retry_loop(save_path: Path, data: bytes): while True: try: save_path.write_bytes(data) break except Exception as e: print('Error while saving checkpoint, retrying in 1 minute: ', e) time.sleep(60) # Ready to train records = [] model.train() with ( train_data_pipe, tqdm(initial=initial_step, total=num_iterations, desc='Training', disable=not accelerator.is_main_process) as pbar, ThreadPoolExecutor(max_workers=1) as save_checkpoint_executor, ): # Get some batches for visualization if accelerator.is_main_process: batches_for_vis: List[Dict[str, torch.Tensor]] = [] num_vis_images = num_vis_images // batch_size_forward * batch_size_forward for _ in range(num_vis_images // batch_size_forward): batch = train_data_pipe.get() batches_for_vis.append(batch) # Visualize GT if vis_every > 0 and accelerator.is_main_process and initial_step == 0: save_dir = Path(workspace).joinpath('vis/gt') for i_batch, batch in enumerate(tqdm(batches_for_vis, desc='Visualize GT', leave=False)): image, gt_depth, gt_normal, gt_intrinsics, info = batch['image'], batch['depth'], batch['normal'], batch['intrinsics'], batch['info'] gt_points = utils3d.pt.depth_map_to_point_map(gt_depth, intrinsics=gt_intrinsics) for i_instance in range(batch['image'].shape[0]): idx = i_batch * batch_size_forward + i_instance image_i = (image[i_instance].numpy().transpose(1, 2, 0) * 255).astype(np.uint8) gt_depth_i = gt_depth[i_instance].numpy() gt_points_i = gt_points[i_instance].numpy() gt_normal_i = gt_normal[i_instance].numpy() save_dir.joinpath(f'{idx:04d}').mkdir(parents=True, exist_ok=True) cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/image.jpg')), cv2.cvtColor(image_i, cv2.COLOR_RGB2BGR)) cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/points.exr')), cv2.cvtColor(gt_points_i, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/depth_vis.png')), cv2.cvtColor(colorize_depth(gt_depth_i), cv2.COLOR_RGB2BGR)) cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/normal.png')), cv2.cvtColor(colorize_normal(gt_normal_i), cv2.COLOR_RGB2BGR)) with save_dir.joinpath(f'{idx:04d}/info.json').open('w') as f: json.dump(info[i_instance], f) # Reset seed to avoid training on the same data when resuming training if seed is not None: set_seed(seed + initial_step, device_specific=True) # Training loop for i_step in range(initial_step, num_iterations): i_accumulate, weight_accumulate = 0, 0 while i_accumulate < gradient_accumulation_steps: # Load batch batch = train_data_pipe.get() image, gt_depth, gt_normal, gt_mask_fin, gt_mask_inf, gt_intrinsics, label_type, is_metric = batch['image'], batch['depth'], batch['normal'], batch['depth_mask_fin'], batch['depth_mask_inf'], batch['intrinsics'], batch['label_type'], batch['is_metric'] image, gt_depth, gt_normal, gt_mask_fin, gt_mask_inf, gt_intrinsics = image.to(device), gt_depth.to(device), gt_normal.to(device), gt_mask_fin.to(device), gt_mask_inf.to(device), gt_intrinsics.to(device) current_batch_size = image.shape[0] if all(label == 'invalid' for label in label_type): continue # NOTE: Skip all-invalid batches to avoid messing up the optimizer. gt_points = utils3d.pt.depth_map_to_point_map(gt_depth, intrinsics=gt_intrinsics) gt_focal = 1 / (1 / gt_intrinsics[..., 0, 0] ** 2 + 1 / gt_intrinsics[..., 1, 1] ** 2) ** 0.5 with accelerator.accumulate(model): # Forward if i_step <= config.get('low_resolution_training_steps', 0): num_tokens = config['model']['num_tokens_range'][0] else: num_tokens = accelerate.utils.broadcast_object_list([random.randint(*config['model']['num_tokens_range'])])[0] with torch.autocast(device_type=accelerator.device.type, dtype=torch.float16, enabled=enable_mixed_precision): output = model(image, num_tokens=num_tokens) pred_points, pred_mask, pred_normal, pred_metric_scale = (output.get(k, None) for k in ['points', 'mask', 'normal', 'metric_scale']) # Compute loss (per instance) loss_list, weight_list = [], [] for i in range(current_batch_size): gt_metric_scale = None loss_dict, weight_dict, misc_dict = {}, {}, {} misc_dict['monitoring'] = monitoring(pred_points[i]) for k, v in config['loss'][label_type[i]].items(): weight_dict[k] = v['weight'] if v['function'] == 'affine_invariant_global_loss': loss_dict[k], misc_dict[k], gt_metric_scale = affine_invariant_global_loss(pred_points[i], gt_points[i], **v['params']) elif v['function'] == 'affine_invariant_local_loss': loss_dict[k], misc_dict[k] = affine_invariant_local_loss(pred_points[i], gt_points[i], gt_focal[i], gt_metric_scale, **v['params']) elif v['function'] == 'normal_loss': loss_dict[k], misc_dict[k] = normal_loss(pred_points[i], gt_points[i]) elif v['function'] == 'edge_loss': loss_dict[k], misc_dict[k] = edge_loss(pred_points[i], gt_points[i]) elif v['function'] == 'normal_map_loss': loss_dict[k], misc_dict[k] = normal_map_loss(pred_normal[i], gt_normal[i]) elif v['function'] == 'mask_bce_loss': loss_dict[k], misc_dict[k] = mask_bce_loss(pred_mask[i], gt_mask_fin[i], gt_mask_inf[i]) elif v['function'] == 'mask_l2_loss': loss_dict[k], misc_dict[k] = mask_l2_loss(pred_mask[i], gt_mask_fin[i], gt_mask_inf[i]) elif v['function'] == 'metric_scale_loss': if is_metric[i] and pred_metric_scale is not None: loss_dict[k], misc_dict[k] = metric_scale_loss(pred_metric_scale[i], gt_metric_scale) else: raise ValueError(f'Undefined loss function: {v["function"]}') weight_dict = {'.'.join(k): v for k, v in flatten_nested_dict(weight_dict).items()} loss_dict = {'.'.join(k): v for k, v in flatten_nested_dict(loss_dict).items()} loss_ = sum([weight_dict[k] * loss_dict[k] for k in loss_dict], start=torch.tensor(0.0, device=device)) loss_list.append(loss_) if torch.isnan(loss_).item(): pbar.write(f'NaN loss in process {accelerator.process_index}') pbar.write(str(loss_dict)) misc_dict = {'.'.join(k): v for k, v in flatten_nested_dict(misc_dict).items()} records.append({ **{k: v.item() for k, v in loss_dict.items()}, **misc_dict, }) loss = sum(loss_list) / len(loss_list) # Backward & update accelerator.backward(loss) if accelerator.sync_gradients: if not enable_mixed_precision and any(torch.isnan(p.grad).any() for p in model.parameters() if p.grad is not None): if accelerator.is_main_process: pbar.write(f'NaN gradients, skip update') optimizer.zero_grad() continue accelerator.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() optimizer.zero_grad() i_accumulate += 1 lr_scheduler.step() # EMA update if enable_ema and accelerator.is_main_process and accelerator.sync_gradients: ema_model.update_parameters(model) # Log metrics if i_step == initial_step or i_step % log_every == 0: records = [key_average(records)] records = accelerator.gather_for_metrics(records, use_gather_object=True) if accelerator.is_main_process: records = key_average(records) if enable_mlflow: try: mlflow.log_metrics(records, step=i_step) except Exception as e: print(f'Error while logging metrics to mlflow: {e}') records = [] # Save model weight checkpoint if accelerator.is_main_process and (i_step % save_every == 0): # NOTE: Writing checkpoint is done in a separate thread to avoid blocking the main process pbar.write(f'Save checkpoint: {i_step:08d}') Path(workspace, 'checkpoint').mkdir(parents=True, exist_ok=True) # Model checkpoint with io.BytesIO() as f: torch.save({ 'model_config': config['model'], 'model': accelerator.unwrap_model(model).state_dict(), }, f) checkpoint_bytes = f.getvalue() save_checkpoint_executor.submit( _write_bytes_retry_loop, Path(workspace, 'checkpoint', f'{i_step:08d}.pt'), checkpoint_bytes ) # Optimizer checkpoint with io.BytesIO() as f: torch.save({ 'model_config': config['model'], 'step': i_step, 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), }, f) checkpoint_bytes = f.getvalue() save_checkpoint_executor.submit( _write_bytes_retry_loop, Path(workspace, 'checkpoint', f'{i_step:08d}_optimizer.pt'), checkpoint_bytes ) # EMA model checkpoint if enable_ema: with io.BytesIO() as f: torch.save({ 'model_config': config['model'], 'model': ema_model.module.state_dict(), }, f) checkpoint_bytes = f.getvalue() save_checkpoint_executor.submit( _write_bytes_retry_loop, Path(workspace, 'checkpoint', f'{i_step:08d}_ema.pt'), checkpoint_bytes ) # Latest checkpoint with io.BytesIO() as f: torch.save({ 'model_config': config['model'], 'step': i_step, }, f) checkpoint_bytes = f.getvalue() save_checkpoint_executor.submit( _write_bytes_retry_loop, Path(workspace, 'checkpoint', 'latest.pt'), checkpoint_bytes ) # Visualize if vis_every > 0 and accelerator.is_main_process and (i_step == initial_step or i_step % vis_every == 0): unwrapped_model = accelerator.unwrap_model(model) save_dir = Path(workspace).joinpath(f'vis/step_{i_step:08d}') save_dir.mkdir(parents=True, exist_ok=True) with torch.inference_mode(): for i_batch, batch in enumerate(tqdm(batches_for_vis, desc=f'Visualize: {i_step:08d}', leave=False)): image, gt_depth, gt_intrinsics = batch['image'], batch['depth'], batch['intrinsics'] image, gt_depth, gt_intrinsics = image.to(device), gt_depth.to(device), gt_intrinsics.to(device) output = unwrapped_model.infer(image) pred_points = output['points'].cpu().numpy() if 'points' in output else None pred_depth = output['depth'].cpu().numpy() if 'depth' in output else None pred_mask = output['mask'].cpu().numpy() if 'mask' in output else None pred_normal = output['normal'].cpu().numpy() if 'normal' in output else None pred_uncertainty = output['uncertainty'].cpu().numpy() if 'uncertainty' in output else None image = (image.cpu().numpy().transpose(0, 2, 3, 1) * 255).astype(np.uint8) for i_instance in range(image.shape[0]): idx = i_batch * batch_size_forward + i_instance save_dir.joinpath(f'{idx:04d}').mkdir(parents=True, exist_ok=True) cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/image.jpg')), cv2.cvtColor(image[i_instance], cv2.COLOR_RGB2BGR)) if pred_points is not None: cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/points.exr')), cv2.cvtColor(pred_points[i_instance], cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) if pred_mask is not None: cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/mask.png')), pred_mask[i_instance] * 255) if pred_depth is not None: cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/depth_vis.png')), cv2.cvtColor(colorize_depth(pred_depth[i_instance], pred_mask[i_instance] if pred_mask is not None else None), cv2.COLOR_RGB2BGR)) if pred_normal is not None: cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/normal_vis.png')), cv2.cvtColor(colorize_normal(pred_normal[i_instance], pred_mask[i_instance] if pred_mask is not None else None), cv2.COLOR_RGB2BGR)) pbar.set_postfix({'loss': loss.item()}, refresh=False) pbar.update(1) if __name__ == '__main__': main()