Whole families supported advances #21

pab1s · Apr 22, 2024 · 727e31f · 727e31f
1 parent f461afc
commit 727e31f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/models/efficientnet.py b/models/efficientnet.py
@@ -24,6 +24,16 @@ def get_efficientnet(model_name, num_classes, pretrained=True):
         model = models.efficientnet_b1(weights=weights)
     elif model_name == "efficientnet_b2":
         model = models.efficientnet_b2(weights=weights)
+    elif model_name == "efficientnet_b3":
+        model = models.efficientnet_b3(weights=weights)
+    elif model_name == "efficientnet_b4":
+        model = models.efficientnet_b4(weights=weights)
+    elif model_name == "efficientnet_b5":
+        model = models.efficientnet_b5(weights=weights)
+    elif model_name == "efficientnet_b6":
+        model = models.efficientnet_b6(weights=weights)
+    elif model_name == "efficientnet_b7":
+        model = models.efficientnet_b7(weights=weights)
     else:
         raise ValueError("Unsupported EfficientNet version")
 

diff --git a/models/resnet.py b/models/resnet.py
@@ -24,6 +24,10 @@ def get_resnet(model_name, num_classes, pretrained=True):
         model = models.resnet34(weights=weights)
     elif model_name == "resnet50":
         model = models.resnet50(weights=weights)
+    elif model_name == "resnet101":
+        model = models.resnet101(weights=weights)
+    elif model_name == "resnet152":
+        model = models.resnet152(weights=weights)
     else:
         raise ValueError("Unsupported ResNet version")