Add default resnet18 from torch to model dict

tsimcne/models/simclr_like.py

@@ -357,6 +357,12 @@ def forward(self, x):
         return self.model.avgpool(feat)
 
 
+def pytorch_resnet(in_channel=3):
+    model = torchvision.models.resnet18(pretrained=False)
+    model.fc = torch.nn.Identity()
+    return model
+
+
 def resnet18(**kwargs):
     return ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
 
@@ -394,7 +400,8 @@ def efficientnet_v2_l(**kwargs):
 
 
 model_dict = {
-    "resnet18": [resnet18, 512],
+    "resnet18": [pytorch_resnet, 512],
+    "resnet18_simclr": [resnet18, 512],
     "resnet34": [resnet34, 512],
     "resnet50": [resnet50, 2048],
     "resnet101": [resnet101, 2048],

tsimcne/tsimcne.py

@@ -20,7 +20,7 @@ def __init__(
         model=None,
         loss="infonce",
         metric=None,
-        backbone="resnet18",
+        backbone="resnet18_simclr",
         projection_head="mlp",
         n_epochs=100,
         batch_size=512,
@@ -160,9 +160,11 @@ class TSimCNE:
         kernel).  Another option is ``"cosine"`` to get the default
         SimCLR loss.
 
-    :param "resnet18" backbone: Backbone to use for the contrastive
-        model.  Defaults to ResNet18.  Other options are
-        ``"resnet50"``, etc. or simply pass in a torch model directly.
+    :param "resnet18_simclr" backbone: Backbone to use for the
+        contrastive model.  Defaults to ResNet18 as defined in the
+        original SimCLR paper (so with a smaller kernel size).  Other
+        options are ``"resnet50"``, etc. or simply pass in a torch
+        model directly.
 
     :param "mlp" projection_head: The projection head that maps from
         the backbone features down to the ``"out_dim"``.  Also accepts
@@ -275,7 +277,7 @@ def __init__(
         model=None,
         loss="infonce",
         metric=None,
-        backbone="resnet18",
+        backbone="resnet18_simclr",
         projection_head="mlp",
         data_transform=None,
         total_epochs=[1000, 50, 450],