()
| 470 | |
| 471 | |
| 472 | def main(): |
| 473 | parser = argparse.ArgumentParser(description="Zero-shot ImageNet evaluation for VTP HuggingFace model") |
| 474 | parser.add_argument("--model_path", type=str, required=True, |
| 475 | help="Path to VTP HuggingFace model directory") |
| 476 | parser.add_argument("--data_path", type=str, required=True, |
| 477 | help="Path to ImageNet validation dataset") |
| 478 | parser.add_argument("--batch_size", type=int, default=128, |
| 479 | help="Batch size for evaluation") |
| 480 | parser.add_argument("--num_workers", type=int, default=8, |
| 481 | help="Number of dataloader workers") |
| 482 | parser.add_argument("--device", type=str, default="cuda:0", |
| 483 | help="Device to use (e.g., cuda:0, cpu)") |
| 484 | parser.add_argument("--precision", type=str, default="fp32", |
| 485 | choices=["fp32", "fp16", "bf16"], |
| 486 | help="Precision for inference") |
| 487 | args = parser.parse_args() |
| 488 | |
| 489 | device = torch.device(args.device) |
| 490 | |
| 491 | print("=" * 60) |
| 492 | print("Zero-Shot ImageNet Evaluation (VTP HuggingFace)") |
| 493 | print("=" * 60) |
| 494 | print(f"Model path: {args.model_path}") |
| 495 | print(f"Data path: {args.data_path}") |
| 496 | print(f"Device: {device}") |
| 497 | print(f"Precision: {args.precision}") |
| 498 | print(f"Batch size: {args.batch_size}") |
| 499 | print() |
| 500 | |
| 501 | # Load model |
| 502 | print("Loading model...") |
| 503 | model = VTPModel.from_pretrained(args.model_path) |
| 504 | model = model.to(device) |
| 505 | model.eval() |
| 506 | |
| 507 | # Get image size from config |
| 508 | image_size = model.config.image_size |
| 509 | context_length = model.config.text_context_length |
| 510 | print(f"Image size: {image_size}") |
| 511 | print(f"Context length: {context_length}") |
| 512 | |
| 513 | # Load tokenizer |
| 514 | print("Loading tokenizer...") |
| 515 | tokenizer = get_tokenizer('ViT-B-32', context_length=context_length) |
| 516 | |
| 517 | # Create dataloader |
| 518 | print(f"Loading ImageNet validation: {args.data_path}") |
| 519 | dataloader = create_imagenet_dataloader( |
| 520 | args.data_path, |
| 521 | image_size=image_size, |
| 522 | batch_size=args.batch_size, |
| 523 | num_workers=args.num_workers, |
| 524 | ) |
| 525 | num_samples = len(dataloader.dataset) |
| 526 | print(f"Number of samples: {num_samples}") |
| 527 | if num_samples != 50000: |
| 528 | print(f" [Warning] Expected 50000 samples for ImageNet validation, got {num_samples}") |
| 529 | print() |
no test coverage detected