()
| 468 | |
| 469 | |
| 470 | def main(): |
| 471 | parser = argparse.ArgumentParser(description="ImageNet reconstruction evaluation for VTP HuggingFace model") |
| 472 | parser.add_argument("--model_path", type=str, required=True, |
| 473 | help="Path to VTP HuggingFace model directory") |
| 474 | parser.add_argument("--data_path", type=str, required=True, |
| 475 | help="Path to ImageNet validation dataset") |
| 476 | parser.add_argument("--output_path", type=str, default="reconstruction_output", |
| 477 | help="Output directory for reconstructed images") |
| 478 | parser.add_argument("--batch_size", type=int, default=32, |
| 479 | help="Batch size per GPU") |
| 480 | parser.add_argument("--num_workers", type=int, default=4, |
| 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="bf16", |
| 485 | choices=["fp32", "fp16", "bf16"], |
| 486 | help="Precision for inference") |
| 487 | parser.add_argument("--no_save_images", action="store_true", |
| 488 | help="Do not save reconstructed images (skip FID calculation)") |
| 489 | parser.add_argument("--use_ddp", action="store_true", |
| 490 | help="Use Distributed Data Parallel for multi-GPU") |
| 491 | parser.add_argument("--local_rank", type=int, default=None, |
| 492 | help="Local rank for DDP (usually set automatically)") |
| 493 | parser.add_argument("--max_samples", type=int, default=None, |
| 494 | help="Maximum number of samples to process (for quick testing)") |
| 495 | args = parser.parse_args() |
| 496 | |
| 497 | # Handle DDP local rank |
| 498 | if args.use_ddp: |
| 499 | if args.local_rank is None: |
| 500 | args.local_rank = int(os.environ.get("LOCAL_RANK", 0)) |
| 501 | os.environ["LOCAL_RANK"] = str(args.local_rank) |
| 502 | |
| 503 | test_reconstruction( |
| 504 | model_path=args.model_path, |
| 505 | data_path=args.data_path, |
| 506 | output_path=args.output_path, |
| 507 | device=args.device, |
| 508 | batch_size=args.batch_size, |
| 509 | precision=args.precision, |
| 510 | save_images=not args.no_save_images, |
| 511 | use_ddp=args.use_ddp, |
| 512 | num_workers=args.num_workers, |
| 513 | max_samples=args.max_samples, |
| 514 | ) |
| 515 | |
| 516 | |
| 517 | if __name__ == "__main__": |
no test coverage detected