def finetune_eval(
    model: Module,
    eval_method: str,
    train_dir: Path,
    val_dir: Path,
    log_dir: Path,
    batch_size_per_device: int,
    num_workers: int,
    accelerator: str,
    devices: int,
    precision: str,
    strategy: str,
    num_classes: int,
) -> Dict[str, float]:
    """Runs fine-tune evaluation on the given model."""
    print_rank_zero("Running fine-tune evaluation...")

    # Setup training data.

    if eval_method == "mae":
        # NOTE: We use transforms from the timm library here as they are the default in MAE
        # and torchvision does not provide all required parameters.
        train_transform = create_transform(
            input_size=224,
            is_training=True,
            auto_augment="rand-m9-mstd0.5-inc1",
            interpolation="bicubic",
            re_prob=0.25,
            re_mode="pixel",
            re_count=1,
            mean=IMAGENET_NORMALIZE["mean"],
            std=IMAGENET_NORMALIZE["std"],
        )
        print_rank_zero("Using MAE training transform.")
    else:
        train_transform = T.Compose(
            [
                T.RandomResizedCrop(224),
                T.RandomHorizontalFlip(),
                T.ToTensor(),
                T.Normalize(
                    mean=IMAGENET_NORMALIZE["mean"], std=IMAGENET_NORMALIZE["std"]
                ),
            ]
        )
        print_rank_zero("Using SimCLR training transform.")

    train_dataset = LightlyDataset(input_dir=str(train_dir), transform=train_transform)
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=batch_size_per_device,
        shuffle=True,
        num_workers=num_workers,
        drop_last=True,
        persistent_workers=True,
    )

    # Setup validation data.
    val_transform = T.Compose(
        [
            T.Resize(256),
            T.CenterCrop(224),
            T.ToTensor(),
            T.Normalize(mean=IMAGENET_NORMALIZE["mean"], std=IMAGENET_NORMALIZE["std"]),
        ]
    )
    val_dataset = LightlyDataset(input_dir=str(val_dir), transform=val_transform)
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=batch_size_per_device,
        shuffle=False,
        num_workers=num_workers,
        persistent_workers=True,
    )

    # Train linear classifier.
    metric_callback = MetricCallback()
    trainer = Trainer(
        max_epochs=100,
        accelerator=accelerator,
        devices=devices,
        callbacks=[
            LearningRateMonitor(),
            DeviceStatsMonitor(),
            metric_callback,
        ],
        logger=TensorBoardLogger(save_dir=str(log_dir), name="finetune_eval"),
        precision=precision,
        strategy=strategy,
        num_sanity_val_steps=0,  # NOTE: save shared memory usage
    )
    if eval_method == "mae":
        classifier = FinetuneClassifierMAE(
            model=model,
            batch_size_per_device=batch_size_per_device,
            feature_dim=model.online_classifier.feature_dim,
            num_classes=num_classes,
        )
        print_rank_zero("Using MAE finetune classifier.")
    else:
        classifier = FinetuneClassifier(
            model=model,
            batch_size_per_device=batch_size_per_device,
            feature_dim=model.online_classifier.feature_dim,
            num_classes=num_classes,
        )
        print_rank_zero("Using SimCLR finetune classifier.")

    trainer.fit(
        model=classifier,
        train_dataloaders=train_dataloader,
        val_dataloaders=val_dataloader,
    )
    metrics_dict: Dict[str, float] = dict()
    for metric in ["val_top1", "val_top5"]:
        print_rank_zero(
            f"max finetune {metric}: {max(metric_callback.val_metrics[metric])}"
        )
        metrics_dict[metric] = max(metric_callback.val_metrics[metric])
    return metrics_dict
Runs fine-tune evaluation on the given model.
benchmarks/imagenet/vitb16/finetune_eval.py