changed to not download pretrained weights for ResNet18

This commit is contained in:
Brian Bjarke Jensen
2024-07-28 21:01:19 +02:00
parent 6c354b23bd
commit e6fdbfe0b0
+2 -2
View File
@@ -9,7 +9,7 @@ class ResNet18Head(nn.Module):
super().__init__()
# copy out parts from ResNet18 with weights
resnet18 = torchvision.models.resnet18(
weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
# weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
)
# save relevant layers
self.conv1 = resnet18.conv1
@@ -21,7 +21,7 @@ class ResNet18Head(nn.Module):
self.layer3 = resnet18.layer3
self.layer4 = resnet18.layer4
self.avgpool = resnet18.avgpool
self.flat = nn.Flatten() # size 512
self.flat = nn.Flatten() # size 512
def forward(self, x):
x = self.conv1(x)