hb

MLRichter · May 2, 2022 · d022015 · d022015
1 parent 63ed330
commit d022015
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/models/efficientnet.py b/models/efficientnet.py
@@ -456,6 +456,22 @@ def efficientnet_b0(*args, **kwargs):
     model.num_classes = ...
     return model
 
+
+@register_model
+def efficientnet_b1(*args, **kwargs):
+    model = _efficientnet_b0(**kwargs)
+    model.name = "EfficentNetB1"
+    model.num_classes = ...
+    return model
+
+
+@register_model
+def efficientnet_b2(*args, **kwargs):
+    model = _efficientnet_b0(**kwargs)
+    model.name = "EfficentNetB2"
+    model.num_classes = ...
+    return model
+
 @register_model
 def efficientnet_b0_eff(*args, **kwargs):
     bneck_conf = partial(MBConvConfig, width_mult=1.0, depth_mult=1.0)