Function create_linear_input

tools/test_linear_probing_hf.py:137–152

Create input for linear classifier from intermediate features.

(x_tokens_list, use_n_blocks: int, use_avgpool: bool)

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

def create_linear_input(x_tokens_list, use_n_blocks: int, use_avgpool: bool) -> torch.Tensor:
    """Create input for linear classifier from intermediate features."""
    intermediate_output = x_tokens_list[-use_n_blocks:]
    output = torch.cat([class_token for _, class_token in intermediate_output], dim=-1)

    if use_avgpool:
        output = torch.cat(
            (
                output,
                torch.mean(intermediate_output[-1][0], dim=1),  # patch tokens
            ),
            dim=-1,
        )
        output = output.reshape(output.shape[0], -1)

    return output.float()


class LinearClassifier(nn.Module):