Upload 12 files

efficientnet_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9223524b7fe0151348f5fd8472beec765c8b57855113b4c4c1a001e53160f6e1
+size 31318643

model.py

@@ -0,0 +1,38 @@
+import torch
+import torchvision
+from torch import nn
+
+def create_effnetb2_model(num_classes: int = 10):
+    """Creates an EfficientNetB2 feature extractor model and transforms.
+
+    Args:
+        num_classes (int, optional): number of classes in the classifier head.
+            Defaults to 3.
+        seed (int, optional): random seed value. Defaults to 42.
+
+    Returns:
+        model (torch.nn.Module): EffNetB2 feature extractor model.
+        transforms (torchvision.transforms): EffNetB2 image transforms.
+    """
+    # Create EffNetB2 pretrained weights, transforms and model
+    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
+    # transforms = weights.transforms()
+    transforms = torchvision.transforms
+    transform = transforms.Compose([transforms.Resize(224),
+                                    transforms.CenterCrop(224),
+                                    transforms.ToTensor(),
+                                    transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])
+
+    model = torchvision.models.efficientnet_b2(weights=weights)
+
+    # Freeze all layers in base model
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # Change classifier head with random seed for reproducibility
+    model.classifier = nn.Sequential(
+        nn.Dropout(p=0.3, inplace=True),
+        nn.Linear(in_features=1408, out_features=num_classes),
+    )
+
+    return model, transform