Skip to content

Commit 567db23

Browse files
committed
Add generic constructor
Add a parameterized constructor for creating any of the EfficientNet variants, i.e. from EfficientNetB0 to EfficientNetB7.
1 parent d740fc8 commit 567db23

File tree

3 files changed

+15
-0
lines changed

3 files changed

+15
-0
lines changed

efficientnet/keras.py

+1
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
from .preprocessing import center_crop_and_resize
55

6+
EfficientNetBn = inject_keras_modules(model.EfficientNetBn)
67
EfficientNetB0 = inject_keras_modules(model.EfficientNetB0)
78
EfficientNetB1 = inject_keras_modules(model.EfficientNetB1)
89
EfficientNetB2 = inject_keras_modules(model.EfficientNetB2)

efficientnet/model.py

+13
Original file line numberDiff line numberDiff line change
@@ -473,6 +473,19 @@ def EfficientNet(width_coefficient,
473473
return model
474474

475475

476+
def EfficientNetBn(n, *args, **kwargs):
477+
return [
478+
EfficientNetB0,
479+
EfficientNetB1,
480+
EfficientNetB2,
481+
EfficientNetB3,
482+
EfficientNetB4,
483+
EfficientNetB5,
484+
EfficientNetB6,
485+
EfficientNetB7,
486+
][n](*args, **kwargs)
487+
488+
476489
def EfficientNetB0(include_top=True,
477490
weights='imagenet',
478491
input_tensor=None,

efficientnet/tfkeras.py

+1
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
from .preprocessing import center_crop_and_resize
55

6+
EfficientNetBn = inject_tfkeras_modules(model.EfficientNetBn)
67
EfficientNetB0 = inject_tfkeras_modules(model.EfficientNetB0)
78
EfficientNetB1 = inject_tfkeras_modules(model.EfficientNetB1)
89
EfficientNetB2 = inject_tfkeras_modules(model.EfficientNetB2)

0 commit comments

Comments
 (0)