Making switching between various GPUS and CPU easy.

ColorizeTrainingArtistic.ipynb

@@ -22,8 +22,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "os.environ['CUDA_VISIBLE_DEVICES']='0' "
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)"
    ]
   },
   {
@@ -32,6 +35,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
     "import fastai\n",
     "from fastai import *\n",
     "from fastai.vision import *\n",
@@ -628,5 +632,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

ColorizeTrainingStable.ipynb

@@ -22,8 +22,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "os.environ['CUDA_VISIBLE_DEVICES']='0' "
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)"
    ]
   },
   {
@@ -32,6 +35,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
     "import fastai\n",
     "from fastai import *\n",
     "from fastai.vision import *\n",
@@ -624,9 +628,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.0"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

ColorizeTrainingVideo.ipynb

@@ -22,8 +22,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "os.environ['CUDA_VISIBLE_DEVICES']='0' "
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)"
    ]
   },
   {
@@ -32,6 +35,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
     "import fastai\n",
     "from fastai import *\n",
     "from fastai.vision import *\n",
@@ -452,9 +456,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.0"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

ColorizeTrainingWandb.ipynb

@@ -34,8 +34,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "os.environ['CUDA_VISIBLE_DEVICES']='0' "
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)"
    ]
   },
   {
@@ -44,6 +47,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
     "import fastai\n",
     "from fastai import *\n",
     "from fastai.vision import *\n",
@@ -844,9 +848,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python (deoldify)",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "deoldify"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -858,9 +862,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.3"
+   "version": "3.7.6"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

ImageColorizer.ipynb

@@ -6,19 +6,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "os.environ['CUDA_VISIBLE_DEVICES']='0' "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "#import numpy as np\n",
-    "#import cv2\n",
-    "#import os"
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)"
    ]
   },
   {
@@ -45,7 +37,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "colorizer = get_image_colorizer(artistic=False)"
+    "colorizer = get_image_colorizer(artistic=True)"
    ]
   },
   {
@@ -148,7 +140,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.6"
+   "version": "3.7.0"
   },
   "toc": {
    "colors": {

ImageColorizerArtisticTests.ipynb

@@ -6,8 +6,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "os.environ['CUDA_VISIBLE_DEVICES']='0' "
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)"
    ]
   },
   {
@@ -3310,5 +3313,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

ImageColorizerColab.ipynb

@@ -56,22 +56,6 @@
     "* Hardware Accelerator = GPU \n"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {},
-    "colab_type": "code",
-    "id": "00_GcC_trpdE"
-   },
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "\n",
-    "if not torch.cuda.is_available():\n",
-    "    print('GPU not available.')"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -114,6 +98,28 @@
     "#◢ Setup"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {},
+    "colab_type": "code",
+    "id": "00_GcC_trpdE"
+   },
+   "outputs": [],
+   "source": [
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "if not torch.cuda.is_available():\n",
+    "    print('GPU not available.')"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -138,9 +144,7 @@
    "outputs": [],
    "source": [
     "import fastai\n",
-    "from deoldify.visualize import *\n",
-    "\n",
-    "torch.backends.cudnn.benchmark = True"
+    "from deoldify.visualize import *"
    ]
   },
   {

ImageColorizerColabStable.ipynb

@@ -56,22 +56,6 @@
     "* Hardware Accelerator = GPU \n"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {},
-    "colab_type": "code",
-    "id": "00_GcC_trpdE"
-   },
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "\n",
-    "if not torch.cuda.is_available():\n",
-    "    print('GPU not available.')"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -114,6 +98,28 @@
     "#◢ Setup"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {},
+    "colab_type": "code",
+    "id": "00_GcC_trpdE"
+   },
+   "outputs": [],
+   "source": [
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "if not torch.cuda.is_available():\n",
+    "    print('GPU not available.')"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

ImageColorizerStableTests.ipynb

@@ -6,8 +6,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "os.environ['CUDA_VISIBLE_DEVICES']='0' "
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)"
    ]
   },
   {
@@ -3297,7 +3300,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.0"
+   "version": "3.7.6"
   },
   "toc": {
    "colors": {
@@ -3325,5 +3328,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

VideoColorizer.ipynb

@@ -6,8 +6,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "os.environ['CUDA_VISIBLE_DEVICES']='0' "
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)"
    ]
   },
   {
@@ -17,8 +20,7 @@
    "outputs": [],
    "source": [
     "from deoldify.visualize import *\n",
-    "plt.style.use('dark_background')\n",
-    "torch.backends.cudnn.benchmark=True"
+    "plt.style.use('dark_background')"
    ]
   },
   {
@@ -172,5 +174,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

VideoColorizerColab.ipynb

@@ -59,20 +59,6 @@
     "* Hardware Accelerator = GPU \n"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {},
-    "colab_type": "code",
-    "id": "00_GcC_trpdE"
-   },
-   "outputs": [],
-   "source": [
-    "from os import path\n",
-    "import torch"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -115,6 +101,30 @@
     "#◢ Setup"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {},
+    "colab_type": "code",
+    "id": "00_GcC_trpdE"
+   },
+   "outputs": [],
+   "source": [
+    "#NOTE:  This must be the first call in order to work properly!\n",
+    "from deoldify import device\n",
+    "from deoldify.device_id import DeviceId\n",
+    "#choices:  CPU, GPU0...GPU7\n",
+    "device.set(device=DeviceId.GPU0)\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "if not torch.cuda.is_available():\n",
+    "    print('GPU not available.')\n",
+    "\n",
+    "from os import path"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

deoldify/__init__.py

@@ -0,0 +1,3 @@
+from deoldify._device import _Device
+
+device = _Device()

deoldify/_device.py

@@ -0,0 +1,31 @@
+import os
+from enum import Enum
+from .device_id import DeviceId
+
+#NOTE:  This must be called first before any torch imports in order to work properly!
+
+class DeviceException(Exception):
+    pass
+
+class _Device:
+    def __init__(self):
+        self.set(DeviceId.CPU)
+
+    def is_gpu(self):
+        ''' Returns `True` if the current device is GPU, `False` otherwise. '''
+        return self.current() is not DeviceID.CPU
+  
+    def current(self):
+        return self._current_device
+
+    def set(self, device:DeviceId):     
+        if device == DeviceId.CPU:
+            os.environ['CUDA_VISIBLE_DEVICES']=''
+        else:
+            os.environ['CUDA_VISIBLE_DEVICES']=str(device.value)
+            import torch
+            torch.backends.cudnn.benchmark=False
+        
+        os.environ['OMP_NUM_THREADS']='1'
+        self._current_device = device    
+        return device

deoldify/device_id.py

@@ -0,0 +1,12 @@
+from enum import IntEnum
+
+class DeviceId(IntEnum):
+    GPU0 = 0,
+    GPU1 = 1,
+    GPU2 = 2,
+    GPU3 = 3,
+    GPU4 = 4,
+    GPU5 = 5,
+    GPU6 = 6,
+    GPU7 = 7,
+    CPU = 99

deoldify/filters.py

@@ -24,6 +24,7 @@ class BaseFilter(IFilter):
     def __init__(self, learn: Learner, stats: tuple = imagenet_stats):
         super().__init__()
         self.learn = learn
+        self.device = next(self.learn.model.parameters()).device
         self.norm, self.denorm = normalize_funcs(*stats)
 
     def _transform(self, image: PilImage) -> PilImage:
@@ -43,10 +44,12 @@ class BaseFilter(IFilter):
     def _model_process(self, orig: PilImage, sz: int) -> PilImage:
         model_image = self._get_model_ready_image(orig, sz)
         x = pil2tensor(model_image, np.float32)
+        x = x.to(self.device)
         x.div_(255)
         x, y = self.norm((x, x), do_x=True)
+
         result = self.learn.pred_batch(
-            ds_type=DatasetType.Valid, batch=(x[None].cuda(), y[None]), reconstruct=True
+            ds_type=DatasetType.Valid, batch=(x[None], y[None]), reconstruct=True
         )
         out = result[0]
         out = self.denorm(out.px, do_x=False)