Function test_reconstruction

tools/test_reconstruction_hf.py:191–467

Test ImageNet reconstruction. Args: model_path: Path to VTP HuggingFace model directory data_path: Path to ImageNet validation dataset output_path: Output directory for reconstructed images device: Device to use (ignored if use_ddp=True) batch_size: Batch

(
    model_path: str,
    data_path: str,
    output_path: str = "reconstruction_output",
    device: str = "cuda:0",
    batch_size: int = 32,
    precision: str = "bf16",
    save_images: bool = True,
    use_ddp: bool = False,
    num_workers: int = 4,
    max_samples: int = None,
)

# ============================================================================

def test_reconstruction(
    model_path: str,
    data_path: str,
    output_path: str = "reconstruction_output",
    device: str = "cuda:0",
    batch_size: int = 32,
    precision: str = "bf16",
    save_images: bool = True,
    use_ddp: bool = False,
    num_workers: int = 4,
    max_samples: int = None,
):
    """Test ImageNet reconstruction.

    Args:
        model_path: Path to VTP HuggingFace model directory
        data_path: Path to ImageNet validation dataset
        output_path: Output directory for reconstructed images
        device: Device to use (ignored if use_ddp=True)
        batch_size: Batch size per GPU
        precision: Precision for inference (fp32, fp16, bf16)
        save_images: Whether to save reconstructed images
        use_ddp: Whether to use Distributed Data Parallel
        num_workers: Number of dataloader workers
        max_samples: Maximum number of samples to process (None for all)
    """
    # Initialize DDP if needed
    if use_ddp:
        if not dist.is_initialized():
            dist.init_process_group(backend='nccl')
        local_rank = dist.get_rank()
        world_size = dist.get_world_size()
        torch.cuda.set_device(local_rank)
        device = torch.device(f'cuda:{local_rank}')
        is_main = (local_rank == 0)
    else:
        local_rank, world_size = 0, 1
        device = torch.device(device)
        is_main = True

    if is_main:
        print("=" * 60)
        print("ImageNet Reconstruction Evaluation (VTP HuggingFace)")
        print("=" * 60)
        print(f"Model path: {model_path}")
        print(f"Data path: {data_path}")
        print(f"Device: {device}" + (f", DDP: {world_size} GPUs" if use_ddp else ""))
        print(f"Precision: {precision}")
        print(f"Batch size: {batch_size}")
        print()

    # Load model
    if is_main:
        print("Loading model...")
    model = VTPModel.from_pretrained(model_path)
    model = model.to(device)
    model.eval()