upload model_builder.py

model_builder.py

@@ -0,0 +1,50 @@
+"""
+contains pytorch model code to instantiate a TinyVGG model.
+"""
+import torch 
+from torch import nn 
+import torchvision
+
+def create_model_baseline_effnetb0(out_feats: int, device: torch.device = None) -> torch.nn.Module:
+    weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT
+    model = torchvision.models.efficientnet_b0(weights=weights).to(device)
+
+    for param in model.features.parameters():
+        param.requires_grad = False
+
+    torch.manual_seed(42)
+    torch.cuda.manual_seed(42)
+
+    # change the output layer 
+    model.classifier = torch.nn.Sequential(
+        torch.nn.Dropout(p=0.2, inplace=True), 
+        torch.nn.Linear(in_features=1280, 
+                        out_features=out_feats, 
+                        bias=True)).to(device)
+    
+    model.name = "effnetb0"
+    print(f"[INFO] created a model {model.name}")
+    
+    return model
+
+def create_model_baseline_effnetb2(out_feats: int, device: torch.device = None) -> torch.nn.Module: 
+    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT 
+    model = torchvision.models.efficientnet_b2(weights=weights).to(device)
+
+    for param in model.features.parameters():
+        param.requires_grad = False 
+
+    torch.manual_seed(42)
+    torch.cuda.manual_seed(42)
+
+    model.classifier = nn.Sequential(
+        nn.Dropout(p=0.3, inplace=True), 
+        nn.Linear(in_features=1408,
+                  out_features=out_feats, 
+                  bias=True)
+    ).to(device)
+
+    model.name = "effnetb2"
+    print(f"[INFO] created a model {model.name}")
+    
+    return model