2024-10-20 00:10:27 +02:00
1 changed files with 6 additions and 2 deletions
@@ -5,11 +5,15 @@ import torchvision


 class ResNet18Head(nn.Module):
-    def __init__(self):
+    def __init__(self, download_resnet_weights: bool = False):
        super().__init__()
        # copy out parts from ResNet18 with weights
+        if download_resnet_weights:
+            weights = torchvision.models.ResNet18_Weights.IMAGENET1K_V1
+        else:
+            weights = None
        resnet18 = torchvision.models.resnet18(
-            # weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
+            weights=weights,
        )
        # save relevant layers
        self.conv1 = resnet18.conv1