Function setup_linear_classifiers

tools/test_linear_probing_hf.py:221–254

Setup linear classifiers with different configurations.

(
    sample_output,
    n_last_blocks_list: Tuple[int, ...],
    learning_rates: Tuple[float, ...],
    batch_size: int,
    num_classes: int = 1000,
    device: torch.device = None,
)

def setup_linear_classifiers(
    sample_output,
    n_last_blocks_list: Tuple[int, ...],
    learning_rates: Tuple[float, ...],
    batch_size: int,
    num_classes: int = 1000,
    device: torch.device = None,
) -> Tuple[AllClassifiers, List[Dict]]:
    """Setup linear classifiers with different configurations."""
    linear_classifiers_dict = nn.ModuleDict()
    optim_param_groups = []

    for n in n_last_blocks_list:
        for avgpool in [True]:
            for _lr in learning_rates:
                lr = scale_lr(_lr, batch_size)
                out_dim = create_linear_input(sample_output, use_n_blocks=n, use_avgpool=avgpool).shape[1]
                linear_classifier = LinearClassifier(
                    out_dim, use_n_blocks=n, use_avgpool=avgpool, num_classes=num_classes
                )
                linear_classifier = linear_classifier.to(device)
                classifier_key = f"classifier_{n}_blocks_avgpool_{avgpool}_lr_{lr:.5f}".replace(".", "_")
                if is_main_process():
                    logger.info(f"Create linear classifier {classifier_key} with input_dim={out_dim}")
                linear_classifiers_dict[classifier_key] = linear_classifier
                optim_param_groups.append({"params": linear_classifier.parameters(), "lr": lr})

    linear_classifiers = AllClassifiers(linear_classifiers_dict)
    if dist.is_initialized():
        linear_classifiers = nn.parallel.DistributedDataParallel(
            linear_classifiers, device_ids=[get_rank() % torch.cuda.device_count()]
        )

    return linear_classifiers, optim_param_groups


def train_one_epoch(