EXP-300

- Fix efficientnet.py docstring

default_params.json

@@ -1,22 +1,22 @@
 {
-  "name": "EfficientNet 3D 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": "FLAIR"
+    "volume_size": [256, 256, 64]
   },
   "model": {
     "name": "efficientnet",
-    "use_aux": true
+    "variant": "B0"
   },
   "train": {
     "lr": 1e-4,
-    "decay_steps": 1000,
+    "decay_steps": 100000,
     "decay_rate": 0.96,
     "epoch": 100
   },
-  "desc": "EfficientNet 3D experiment with 256x256x64 voxel size, with aux out."
+  "desc": "EfficientNetB0 3D experiment with 256x256x64 voxel size."
 }

models/efficientnet.py

@@ -2,7 +2,17 @@
 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"""
+    """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,