diff --git a/default_params.json b/default_params.json index 4e4d482..c18637b 100644 --- a/default_params.json +++ b/default_params.json @@ -1,22 +1,22 @@ { - "name": "Ensemble Learning 256 with aux", + "name": "EfficientNetB0 3D 256", "seed": 7171, + "ensemble": true, "data": { "dirname": "data/orig", "val_size": 0.2, "batch_size": 1, - "volume_size": [256, 256, 64], - "seq_type": "ALL" + "volume_size": [256, 256, 64] }, "model": { - "name": "ensemble_model", - "use_aux": true + "name": "efficientnet", + "variant": "B0" }, "train": { "lr": 1e-4, - "decay_steps": 1000, + "decay_steps": 100000, "decay_rate": 0.96, "epoch": 100 }, - "desc": "Ensemble learning experiment with 256x256x64 voxel size, with aux out." + "desc": "EfficientNetB0 3D experiment with 256x256x64 voxel size." } diff --git a/models/__init__.py b/models/__init__.py index 34ec76e..dfce67c 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -9,4 +9,5 @@ def find_model(model_name): # Import exposed model functions from models.baseline import baseline +from models.efficientnet import efficientnet from models.ensemble import ensemble_model diff --git a/models/efficientnet.py b/models/efficientnet.py new file mode 100644 index 0000000..6eeee57 --- /dev/null +++ b/models/efficientnet.py @@ -0,0 +1,41 @@ +from tensorflow import keras +import efficientnet_3D.tfkeras as efn + +def efficientnet(input_shape, n_class=2, variant='B0', **kwargs): + """EfficientNet model from the paper https://arxiv.org/abs/1905.11946 + + Args: + input_shape: Shape of the input tensor, not including the `batch_size` dimension. + n_class: The number of class to be predicted. + variant: The variant of efficient net. Can be `B0`~`B7` for both 2D and 3D. + `L2` is available for 3D version only. + + Returns: + A `keras.Model` instance. + """ + default_effnet_params = { + 'include_top': True, + 'weights': None, + 'input_shape': input_shape, + 'pooling': None, + 'classes': n_class, + } + try: + if len(input_shape) == 4: + effnet_layer = getattr(efn, f'EfficientNet{variant}')(**default_effnet_params) + modal = '3D' + elif len(input_shape) == 3: + effnet_layer = getattr(keras.applications, f'EfficientNet{variant}')( + classifier_activation='softmax', **default_effnet_params + ) + modal = '2D' + else: + raise ValueError('input_shape is expected as an array ranked 3 or 4') + except AttributeError: + print(f'No EfficientNet variation found for {variant}') + raise SystemExit + + inputs = keras.Input(input_shape) + outputs = effnet_layer(inputs) + + return keras.Model(inputs, outputs, name=f'efficientnet_{modal}_{variant}') diff --git a/setup/requirements.txt b/setup/requirements.txt index 7b9cc5c..878eafb 100644 --- a/setup/requirements.txt +++ b/setup/requirements.txt @@ -1,5 +1,7 @@ +efficientnet-3D==1.0.1 imageio==2.9.0 jupyter +Keras-Applications==1.0.8 matplotlib==3.3.2 nibabel==3.2.1 pandas==1.1.4