Question: Is features[0][:, 0] in model.py intended to be the CLS token?

[houxuxing]

Summary

In src/model.py, the ViTBackboneNet class extracts what the comment describes as the CLS token:

features = self.backbone(x)
cls_token = features[0][:, 0]  # Shape: [batch_size, 768]

I would like to clarify whether features[0][:, 0] is indeed the CLS token or actually the first spatial patch token.

Verification

The MONAI ViT implementation (tested with monai==1.3.2 as specified in requirements.txt) only creates a cls_token when classification=True. In model.py, ViT is instantiated without classification=True:

self.backbone = ViT(
    in_channels=1,
    img_size=(96,96,96),
    patch_size=(16, 16, 16),
    hidden_size=768,
    mlp_dim=3072,
    num_layers=12,
    num_heads=12,
    save_attn=True,
)

Running the above configuration gives:

out = self.backbone(torch.randn(1, 1, 96, 96, 96))
print(out[0].shape)                         # torch.Size([1, 216, 768])
print(hasattr(self.backbone, "cls_token"))  # False

Since 216 = (96/16)^3 equals the number of patch tokens, there is no CLS token prepended to the sequence, and features[0][:, 0] corresponds to the first patch token.

Question

Was the SimCLR pretraining also performed without a CLS token (i.e., classification=False)? If so, the downstream code is consistent and the issue is mainly the misleading comment. If pretraining used classification=True, then the cls_token weights would not be loaded by the current downstream code.

Suggested improvement

To avoid ambiguity, consider updating the comment or switching to a patch-aggregation strategy such as mean pooling:

features = features[0].mean(dim=1)  # (batch, 768)

Thanks for open-sourcing BrainIAC!