added mixed precision

README.md

@@ -8,4 +8,14 @@ app_file: app.py
 pinned: false
 ---
 
-# enhance-me
+# Enhance Me
+
+A unified platform for image enhancement.
+
+## Usage
+
+### Train using Docker
+
+- Build image using `docker build -t enhance-image .`
+
+- Run notebook using `docker run -it --gpus all -p 8888:8888 -v $(pwd):/usr/src/enhance-me enhance-image`

enhance_me/zero_dce/zero_dce.py

@@ -5,7 +5,7 @@ from datetime import datetime
 
 import tensorflow as tf
 from tensorflow import keras
-from tensorflow.keras import optimizers, Model
+from tensorflow.keras import optimizers, mixed_precision, Model
 from wandb.keras import WandbCallback
 
 from .dce_net import build_dce_net
@@ -20,9 +20,12 @@ from ..commons import download_lol_dataset, init_wandb
 
 
 class ZeroDCE(Model):
-    def __init__(self, experiment_name=None, wandb_api_key=None, **kwargs):
+    def __init__(self, experiment_name=None, wandb_api_key=None, use_mixed_precision: bool = False, **kwargs):
         super(ZeroDCE, self).__init__(**kwargs)
         self.experiment_name = experiment_name
+        if use_mixed_precision:
+            policy = mixed_precision.Policy('mixed_float16')
+            mixed_precision.set_global_policy(policy)
         if wandb_api_key is not None:
             init_wandb("zero-dce", experiment_name, wandb_api_key)
             self.using_wandb = True

notebooks/enhance_me_train.ipynb

@@ -187,15 +187,16 @@
    "source": [
     "# @title Zero-DCE Train Configs\n",
     "\n",
-    "experiment_name = \"lol_dataset_256\"  # @param {type:\"string\"}\n",
+    "experiment_name = \"lol_dataset_256_resize\"  # @param {type:\"string\"}\n",
     "image_size = 256  # @param {type:\"integer\"}\n",
     "dataset_label = \"lol\"  # @param [\"lol\"]\n",
-    "apply_resize = False  # @param {type:\"boolean\"}\n",
+    "use_mixed_precision = False  # @param {type:\"boolean\"}\n",
+    "apply_resize = True  # @param {type:\"boolean\"}\n",
     "apply_random_horizontal_flip = True  # @param {type:\"boolean\"}\n",
     "apply_random_vertical_flip = True  # @param {type:\"boolean\"}\n",
     "apply_random_rotation = True  # @param {type:\"boolean\"}\n",
     "use_mixed_precision = False  # @param {type:\"boolean\"}\n",
-    "wandb_api_key = \"\"  # @param {type:\"string\"}\n",
+    "wandb_api_key = \"8d7149fe07496df2aaab8e9856a6ed8564e2a644\"  # @param {type:\"string\"}\n",
     "val_split = 0.1  # @param {type:\"slider\", min:0.1, max:1.0, step:0.1}\n",
     "batch_size = 16  # @param {type:\"integer\"}\n",
     "learning_rate = 1e-4  # @param {type:\"number\"}\n",
@@ -211,7 +212,8 @@
    "source": [
     "zero_dce = ZeroDCE(\n",
     "    experiment_name=experiment_name,\n",
-    "    wandb_api_key=None if wandb_api_key == \"\" else wandb_api_key\n",
+    "    wandb_api_key=None if wandb_api_key == \"\" else wandb_api_key,\n",
+    "    use_mixed_precision=use_mixed_precision\n",
     ")"
    ]
   },
@@ -263,6 +265,13 @@
     "        (18, 18),\n",
     "    )"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {