()
| 25 | |
| 26 | |
| 27 | def main(): |
| 28 | set_seed(42) |
| 29 | os.environ["WANDB__SERVICE_WAIT"] = "300" |
| 30 | |
| 31 | opt = tyro.cli(AllConfigs) |
| 32 | |
| 33 | torch.set_float32_matmul_precision('high') |
| 34 | |
| 35 | accelerator = Accelerator( |
| 36 | mixed_precision=opt.mixed_precision, |
| 37 | gradient_accumulation_steps=opt.gradient_accumulation_steps, |
| 38 | ) |
| 39 | |
| 40 | dataset_map = { |
| 41 | "co3d": Co3DDataset, |
| 42 | "re10k": Re10kDataset, |
| 43 | "davis": DavisDataset, |
| 44 | 'vos': VOSDataset, |
| 45 | 'combined': CombinedDataset |
| 46 | } |
| 47 | |
| 48 | for key, Dataset in dataset_map.items(): |
| 49 | if key in opt.root_path.lower(): |
| 50 | dataset_nm = key |
| 51 | print(f"Loading dataset: {key}") |
| 52 | break |
| 53 | else: |
| 54 | raise ValueError(f"Dataset {opt.root_path} not supported") |
| 55 | |
| 56 | # Warmup parameters |
| 57 | # initial_batch_size = 8 |
| 58 | initial_batch_size = opt.batch_size |
| 59 | target_batch_size = opt.batch_size |
| 60 | warmup_epochs = 15 |
| 61 | current_batch_size = initial_batch_size |
| 62 | |
| 63 | train_dataset = Dataset(opt=opt, shuffle=True, training=True) |
| 64 | train_dataloader = torch.utils.data.DataLoader( |
| 65 | train_dataset, |
| 66 | batch_size = current_batch_size, |
| 67 | num_workers=opt.num_workers, |
| 68 | pin_memory=True, |
| 69 | shuffle=not isinstance(train_dataset, torch.utils.data.IterableDataset), |
| 70 | drop_last=True, |
| 71 | ) |
| 72 | |
| 73 | test_dataset = Dataset(opt=opt, shuffle=True, training=False) |
| 74 | test_dataloader = torch.utils.data.DataLoader( |
| 75 | test_dataset, |
| 76 | batch_size=opt.batch_size * 2, |
| 77 | num_workers=opt.num_workers, |
| 78 | pin_memory=True, |
| 79 | drop_last=False, |
| 80 | ) |
| 81 | |
| 82 | model = StaticEncoder(opt) |
| 83 | |
| 84 | optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=opt.lr, weight_decay=0.05, betas=(0.9, 0.95)) |
no test coverage detected