added option to download resnet18 weights

2024-07-28 21:19:47 +02:00
parent 96ccdb5c90
commit 6541a47668
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