()
| 556 | |
| 557 | |
| 558 | def main(): |
| 559 | parser = argparse.ArgumentParser(description="Linear probing evaluation for VTP HuggingFace model") |
| 560 | parser.add_argument("--model_path", type=str, required=True, |
| 561 | help="Path to VTP HuggingFace model directory") |
| 562 | parser.add_argument("--imagenet_root", type=str, required=True, |
| 563 | help="Path to ImageNet dataset root (with train/ and val/ subdirs)") |
| 564 | parser.add_argument("--output_dir", type=str, default="./linear_probing_results", |
| 565 | help="Output directory for results") |
| 566 | parser.add_argument("--batch_size", type=int, default=128, |
| 567 | help="Batch size per GPU") |
| 568 | parser.add_argument("--epochs", type=int, default=10, |
| 569 | help="Number of training epochs") |
| 570 | parser.add_argument("--epoch_length", type=int, default=1250, |
| 571 | help="Number of iterations per epoch") |
| 572 | parser.add_argument("--num_workers", type=int, default=8, |
| 573 | help="Number of dataloader workers") |
| 574 | parser.add_argument("--device", type=str, default="cuda:0", |
| 575 | help="Device to use (e.g., cuda:0)") |
| 576 | parser.add_argument("--precision", type=str, default="bf16", |
| 577 | choices=["fp32", "fp16", "bf16"], |
| 578 | help="Precision for inference") |
| 579 | parser.add_argument("--use_ddp", action="store_true", |
| 580 | help="Use Distributed Data Parallel for multi-GPU") |
| 581 | parser.add_argument("--local_rank", type=int, default=None, |
| 582 | help="Local rank for DDP (usually set automatically)") |
| 583 | args = parser.parse_args() |
| 584 | |
| 585 | if args.use_ddp: |
| 586 | if args.local_rank is None: |
| 587 | args.local_rank = int(os.environ.get("LOCAL_RANK", 0)) |
| 588 | os.environ["LOCAL_RANK"] = str(args.local_rank) |
| 589 | |
| 590 | test_linear_probing( |
| 591 | model_path=args.model_path, |
| 592 | imagenet_root=args.imagenet_root, |
| 593 | output_dir=args.output_dir, |
| 594 | batch_size=args.batch_size, |
| 595 | epochs=args.epochs, |
| 596 | epoch_length=args.epoch_length, |
| 597 | precision=args.precision, |
| 598 | use_ddp=args.use_ddp, |
| 599 | device=args.device, |
| 600 | num_workers=args.num_workers, |
| 601 | ) |
| 602 | |
| 603 | |
| 604 | if __name__ == "__main__": |
no test coverage detected