brian/visual_critical_discourse_analysis
1
0
Fork 0
Code Issues 11 Pull Requests Actions Projects 1 Releases Activity

added option to download resnet18 weights

Browse Source
This commit is contained in:
Brian Bjarke Jensen
2024-07-28 21:19:47 +02:00
parent 96ccdb5c90
commit 6541a47668
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/src/models/resnet18_head.py
+6 -2
View File
@@ -5,11 +5,15 @@ import torchvision
class ResNet18Head(nn.Module): class ResNet18Head(nn.Module):
def __init__(self): def __init__(self, download_resnet_weights: bool = False):
super().__init__() super().__init__()
# copy out parts from ResNet18 with weights # 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( resnet18 = torchvision.models.resnet18(
# weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1, weights=weights,
) )
# save relevant layers # save relevant layers
self.conv1 = resnet18.conv1 self.conv1 = resnet18.conv1