Run linear probing evaluation. Args: model_path: Path to VTP HuggingFace model directory imagenet_root: Path to ImageNet dataset root (with train/ and val/ subdirs) output_dir: Output directory for results batch_size: Batch size per GPU epochs: Number of
(
model_path: str,
imagenet_root: str,
output_dir: str = "./linear_probing_results",
batch_size: int = 128,
epochs: int = 10,
epoch_length: int = 1250,
n_last_blocks_list: Tuple[int, ...] = (1, 4),
learning_rates: Tuple[float, ...] = DEFAULT_LEARNING_RATES,
precision: str = "bf16",
use_ddp: bool = False,
device: str = "cuda:0",
num_workers: int = 8,
)
| 350 | # ============================================================================ |
| 351 | |
| 352 | def test_linear_probing( |
| 353 | model_path: str, |
| 354 | imagenet_root: str, |
| 355 | output_dir: str = "./linear_probing_results", |
| 356 | batch_size: int = 128, |
| 357 | epochs: int = 10, |
| 358 | epoch_length: int = 1250, |
| 359 | n_last_blocks_list: Tuple[int, ...] = (1, 4), |
| 360 | learning_rates: Tuple[float, ...] = DEFAULT_LEARNING_RATES, |
| 361 | precision: str = "bf16", |
| 362 | use_ddp: bool = False, |
| 363 | device: str = "cuda:0", |
| 364 | num_workers: int = 8, |
| 365 | ): |
| 366 | """Run linear probing evaluation. |
| 367 | |
| 368 | Args: |
| 369 | model_path: Path to VTP HuggingFace model directory |
| 370 | imagenet_root: Path to ImageNet dataset root (with train/ and val/ subdirs) |
| 371 | output_dir: Output directory for results |
| 372 | batch_size: Batch size per GPU |
| 373 | epochs: Number of training epochs |
| 374 | epoch_length: Number of iterations per epoch |
| 375 | n_last_blocks_list: Number of last blocks to use for features |
| 376 | learning_rates: Learning rates to sweep |
| 377 | precision: Precision for inference (fp32, fp16, bf16) |
| 378 | use_ddp: Whether to use Distributed Data Parallel |
| 379 | device: Device to use (ignored if use_ddp) |
| 380 | num_workers: Number of dataloader workers |
| 381 | """ |
| 382 | # Initialize DDP if needed |
| 383 | if use_ddp: |
| 384 | if not dist.is_initialized(): |
| 385 | dist.init_process_group(backend='nccl') |
| 386 | local_rank = get_rank() |
| 387 | torch.cuda.set_device(local_rank % torch.cuda.device_count()) |
| 388 | device = torch.device(f'cuda:{local_rank % torch.cuda.device_count()}') |
| 389 | else: |
| 390 | device = torch.device(device) |
| 391 | |
| 392 | if is_main_process(): |
| 393 | print("=" * 60) |
| 394 | print("Linear Probing Evaluation (VTP HuggingFace)") |
| 395 | print("=" * 60) |
| 396 | print(f"Model path: {model_path}") |
| 397 | print(f"ImageNet root: {imagenet_root}") |
| 398 | print(f"Output dir: {output_dir}") |
| 399 | print(f"Device: {device}" + (f", DDP: {get_world_size()} GPUs" if use_ddp else "")) |
| 400 | print(f"Precision: {precision}") |
| 401 | print(f"Batch size: {batch_size}") |
| 402 | print(f"Epochs: {epochs}, Epoch length: {epoch_length}") |
| 403 | print() |
| 404 | |
| 405 | os.makedirs(output_dir, exist_ok=True) |
| 406 | cudnn.benchmark = True |
| 407 | |
| 408 | # Load model |
| 409 | if is_main_process(): |
no test coverage detected