Function main

train_decoder.py:25–390

()

import warnings

def main():
    set_seed(42)
    os.environ["WANDB__SERVICE_WAIT"] = "300"

    opt = tyro.cli(AllConfigs)

    # Load dataset based on root_path
    dataset_map = {
        "co3d": Co3DDataset,
        "re10k": Re10kDataset, 
        "davis": DavisDataset,
        'vos': VOSDataset,
        'combined': CombinedDataset
    }

    for key, Dataset in dataset_map.items():
        if key in opt.root_path.lower():
            dataset_nm = key
            print(f"Loading dataset: {key}")
            break
    else:
        raise ValueError(f"Dataset {opt.root_path} not supported")
    
    # Batch size management
    initial_batch_size = opt.batch_size
    target_batch_size = opt.batch_size
    warmup_epochs = 10
    current_batch_size = initial_batch_size

    torch.set_float32_matmul_precision('high')
    
    accelerator = Accelerator(
        mixed_precision=opt.mixed_precision,
        gradient_accumulation_steps=opt.gradient_accumulation_steps,
    )

    train_dataset = Dataset(opt=opt, shuffle=True, training=True)
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=current_batch_size,
        num_workers=opt.num_workers,
        pin_memory=True,
        shuffle=not isinstance(train_dataset, torch.utils.data.IterableDataset),
        drop_last=True,
    )

    test_dataset = Dataset(opt=opt, shuffle=True, training=False)
    test_dataloader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=opt.batch_size * 2,
        shuffle=not isinstance(test_dataset, torch.utils.data.IterableDataset),
        num_workers=0,
        pin_memory=True,
        drop_last=False,
    )
    
    model = SplatModel(opt)