DM-Diaz commited on
Commit
db2b22f
·
verified ·
1 Parent(s): c000dfb

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +82 -1
README.md CHANGED
@@ -103,6 +103,77 @@ The stored epoch is zero-indexed, so `epoch = 99` corresponds to the completion
103
 
104
  The checkpoint includes training state such as optimizer and scheduler information in addition to model parameters.
105
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
106
  ### Comparative Evaluation Results
107
 
108
  The table below reproduces the summary metrics reported in the associated paper across all VEDB-trained conditions and **reference models**. **Rows corresponding to this repository's ImageNet-1K checkpoint are bolded.**
@@ -177,4 +248,14 @@ If you use this checkpoint or representations derived from it in academic work,
177
  address = {New York, NY, USA},
178
  year = {2026},
179
  doi = {10.32470/0416gfsq}
180
- }
 
 
 
 
 
 
 
 
 
 
 
103
 
104
  The checkpoint includes training state such as optimizer and scheduler information in addition to model parameters.
105
 
106
+ ## Loading the Checkpoint
107
+
108
+ The checkpoint can be loaded by reconstructing the ResNet-18 backbone and SimCLR projection head used during training.
109
+
110
+ ```python
111
+ import torch
112
+ import torch.nn as nn
113
+ import torchvision
114
+ from lightly.models.modules import heads
115
+
116
+
117
+ class SimCLRResNet18(nn.Module):
118
+ def __init__(self):
119
+ super().__init__()
120
+
121
+ resnet = torchvision.models.resnet18(weights=None)
122
+ feature_dim = resnet.fc.in_features # 512
123
+
124
+ # Remove the classification head
125
+ self.backbone = nn.Sequential(
126
+ *list(resnet.children())[:-1]
127
+ )
128
+
129
+ # SimCLR projection head: 512 -> 512 -> 128
130
+ self.projection_head = heads.SimCLRProjectionHead(
131
+ feature_dim,
132
+ feature_dim,
133
+ 128,
134
+ )
135
+
136
+ def forward(self, x):
137
+ features = self.backbone(x).flatten(start_dim=1)
138
+ projections = self.projection_head(features)
139
+ return projections
140
+
141
+
142
+ checkpoint = torch.load(
143
+ "checkpoint_100-resnet18-simclr-imagenet1k.ckpt",
144
+ map_location="cpu",
145
+ weights_only=False,
146
+ )
147
+
148
+ model = SimCLRResNet18()
149
+ model.load_state_dict(checkpoint["state_dict"], strict=True)
150
+ model.eval()
151
+ ```
152
+
153
+ ### Extracting Backbone Features
154
+
155
+ For most downstream applications, the 512-dimensional ResNet-18 representation can be extracted without using the SimCLR projection head:
156
+
157
+ ```python
158
+ with torch.no_grad():
159
+ features = model.backbone(images).flatten(start_dim=1)
160
+
161
+ print(features.shape)
162
+ # [batch_size, 512]
163
+ ```
164
+
165
+ The 128-dimensional SimCLR projection can instead be obtained with:
166
+
167
+ ```python
168
+ with torch.no_grad():
169
+ projections = model(images)
170
+
171
+ print(projections.shape)
172
+ # [batch_size, 128]
173
+ ```
174
+
175
+ Input tensors should have shape `[batch_size, 3, 224, 224]`.
176
+
177
  ### Comparative Evaluation Results
178
 
179
  The table below reproduces the summary metrics reported in the associated paper across all VEDB-trained conditions and **reference models**. **Rows corresponding to this repository's ImageNet-1K checkpoint are bolded.**
 
248
  address = {New York, NY, USA},
249
  year = {2026},
250
  doi = {10.32470/0416gfsq}
251
+ }
252
+ ```
253
+
254
+ **Proceedings:** [Diaz & Henderson (2026)](https://doi.org/10.32470/0416gfsq)<br>
255
+ **Preprint:** [arXiv:2607.19316](https://arxiv.org/abs/2607.19316)
256
+
257
+ ## License
258
+
259
+ The released checkpoint and repository materials are provided under the **Apache License 2.0**.
260
+
261
+ The ImageNet-1K training dataset and third-party software used to produce the model remain subject to their respective licenses, access requirements, and terms of use.