From b13ba7f2ecda4116faa44ee2fd4bcb884d6e46c7 Mon Sep 17 00:00:00 2001
From: magnusbarata <barata.magnuz@gmail.com>
Date: Thu, 30 Sep 2021 03:44:22 +0900
Subject: [PATCH] EXP-300 - Fix efficientnet.py docstring

---
 default_params.json    | 12 ++++++------
 models/efficientnet.py | 12 +++++++++++-
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/default_params.json b/default_params.json
index 6d752e8..c18637b 100644
--- a/default_params.json
+++ b/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."
 }
diff --git a/models/efficientnet.py b/models/efficientnet.py
index e644a3d..6eeee57 100644
--- a/models/efficientnet.py
+++ b/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,