Experimental ONNX runner

Author: tayden

kelp_o_matic/hann.py

@@ -1,162 +1,151 @@
 import math
-from abc import ABCMeta, abstractmethod
+from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Union
 
 import numpy as np
 import rasterio
-import torch
 from rasterio.windows import Window
 
 
 # Implementation of paper:
 # https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0229839#pone.0229839.ref007
 
 
-class Kernel(torch.nn.Module, metaclass=ABCMeta):
-    def __init__(
-        self, size: int = 512, device: torch.device.type = torch.device("cpu")
-    ):
+class Kernel(ABC):
+    def __init__(self, size: int = 512):
         super().__init__()
         self.size = size
-        self.wi = self._init_wi(size, device)
-        self.wj = self.wi.clone()
+        self.wi = self._init_wi(size)
+        self.wj = self.wi.copy()
 
     @staticmethod
     @abstractmethod
-    def _init_wi(size: int, device: torch.device.type) -> torch.Tensor:
+    def _init_wi(size: int) -> np.ndarray:
         raise NotImplementedError
 
     def get_kernel(
-        self,
-        top: bool = False,
-        bottom: bool = False,
-        left: bool = False,
-        right: bool = False,
-    ) -> torch.Tensor:
-        wi, wj = self.wi.clone(), self.wj.clone()
+            self,
+            top: bool = False,
+            bottom: bool = False,
+            left: bool = False,
+            right: bool = False,
+    ) -> "np.ndarray":
+        wi, wj = self.wi.copy(), self.wj.copy()
 
         if top:
             wi[: self.size // 2] = 1
         if bottom:
-            wi[self.size // 2 :] = 1
+            wi[self.size // 2:] = 1
 
         if left:
             wj[: self.size // 2] = 1
         if right:
-            wj[self.size // 2 :] = 1
+            wj[self.size // 2:] = 1
 
-        return wi.unsqueeze(1) @ wj.unsqueeze(0)
+        return np.expand_dims(wi, 1) @ np.expand_dims(wj, 0)
 
     def forward(
-        self,
-        x: torch.Tensor,
-        top: bool = False,
-        bottom: bool = False,
-        left: bool = False,
-        right: bool = False,
-    ) -> torch.Tensor:
+            self,
+            x: "np.ndarray",
+            top: bool = False,
+            bottom: bool = False,
+            left: bool = False,
+            right: bool = False,
+    ) -> np.ndarray:
         kernel = self.get_kernel(top=top, bottom=bottom, left=left, right=right)
-        return torch.mul(x, kernel)
+        return np.multiply(x, kernel)
 
 
 class HannKernel(Kernel):
     @staticmethod
-    def _init_wi(size: int, device: torch.device.type) -> torch.Tensor:
-        i = torch.arange(0, size, device=device)
-        return (1 - ((2 * np.pi * i) / (size - 1)).cos()) / 2
+    def _init_wi(size: int) -> np.ndarray:
+        i = np.arange(0, size)
+        return (1 - np.cos(((2 * np.pi * i) / (size - 1)))) / 2
 
 
 class BartlettHannKernel(Kernel):
     @staticmethod
-    def _init_wi(size: int, device: torch.device.type) -> torch.Tensor:
+    def _init_wi(size: int) -> np.ndarray:
         # Follows original paper:
         # Ha YH, Pearce JA. A new window and comparison to standard windows.
         # IEEE Transactions on Acoustics, Speech, and Signal Processing.
         # 1989;37(2):298–301.
-        i = torch.arange(0, size, device=device)
+        i = np.arange(0, size)
         return (
-            0.62
-            - 0.48 * (i / size - 1 / 2).abs()
-            + 0.38 * (2 * np.pi * (i / size - 1 / 2).abs()).cos()
+                0.62
+                - 0.48 * np.abs((i / size - 1 / 2))
+                + 0.38 * np.cos((2 * np.pi * np.abs((i / size - 1 / 2))))
         )
 
 
 class TriangularKernel(Kernel):
     @staticmethod
-    def _init_wi(size: int, device: torch.device.type) -> torch.Tensor:
-        i = torch.arange(0, size, device=device)
-        return 1 - (2 * i / size - 1).abs()
+    def _init_wi(size: int) -> np.ndarray:
+        i = np.arange(0, size)
+        return 1 - np.abs((2 * i / size - 1))
 
 
 class BlackmanKernel(Kernel):
     @staticmethod
-    def _init_wi(size: int, device: torch.device.type) -> torch.Tensor:
-        i = torch.arange(0, size, device=device)
+    def _init_wi(size: int) -> np.ndarray:
+        i = np.arange(0, size)
         return (
-            0.42
-            - 0.5 * (2 * np.pi * i / size).cos()
-            + 0.08 * (4 * np.pi * i / size).cos()
+                0.42
+                - 0.5 * np.cos((2 * np.pi * i / size))
+                + 0.08 * np.cos((4 * np.pi * i / size))
         )
 
 
-class TorchMemoryRegister(object):
+class NumpyMemoryRegister(object):
     def __init__(
-        self,
-        image_path: Union[str, Path],
-        reg_depth: int,
-        window_size: int,
-        device: torch.device.type,
+            self,
+            image_path: Union[str, Path],
+            reg_depth: int,
+            window_size: int,
     ):
         super().__init__()
         self.image_path = Path(image_path)
         self.n = reg_depth
         self.ws = window_size
         self.hws = window_size // 2
-        self.device = device
 
         # Copy metadata from img
         with rasterio.open(str(image_path), "r") as src:
             src_width = src.width
 
         self.height = self.ws
         self.width = (math.ceil(src_width / self.ws) * self.ws) + self.hws
-        self.register = torch.zeros(
-            (self.n, self.height, self.width), device=self.device
-        )
+        self.register = np.zeros((self.n, self.height, self.width))
 
     @property
     def _zero_chip(self):
-        return torch.zeros(
-            (self.n, self.hws, self.hws), dtype=torch.float, device=self.device
-        )
+        return np.zeros((self.n, self.hws, self.hws), dtype=float)
 
-    def step(self, new_logits: torch.Tensor, img_window: Window):
+    def step(self, new_logits: "np.ndarray", img_window: Window):
         # 1. Read data from the registry to update with the new logits
         # |a|b| |
         # |c|d| |
-        with torch.no_grad():
-            logits_abcd = self.register[
-                :, :, img_window.col_off : img_window.col_off + self.ws
-            ].clone()
-            logits_abcd += new_logits
+        logits_abcd = self.register[:, :,
+                      img_window.col_off: img_window.col_off + self.ws].copy()
+        logits_abcd += new_logits
 
         # Update the registry and pop information-complete data
         # |c|b| | + pop a
         # |0|d| |
         logits_a = logits_abcd[:, : self.hws, : self.hws]
-        logits_c = logits_abcd[:, self.hws :, : self.hws]
-        logits_c0 = torch.concat([logits_c, self._zero_chip], dim=1)
-        logits_bd = logits_abcd[:, :, self.hws :]
+        logits_c = logits_abcd[:, self.hws:, : self.hws]
+        logits_c0 = np.concatenate([logits_c, self._zero_chip], axis=1)
+        logits_bd = logits_abcd[:, :, self.hws:]
 
         # write c0
         self.register[
-            :, :, img_window.col_off : img_window.col_off + self.hws
+        :, :, img_window.col_off: img_window.col_off + self.hws
         ] = logits_c0
 
         # write bd
         col_off_bd = img_window.col_off + self.hws
-        self.register[:, :, col_off_bd : col_off_bd + self.hws] = logits_bd
+        self.register[:, :, col_off_bd: col_off_bd + self.hws] = logits_bd
 
         # Return the information-complete predictions
         preds_win = Window(
@@ -165,6 +154,9 @@ def step(self, new_logits: torch.Tensor, img_window: Window):
             height=min(self.hws, img_window.height),
             width=min(self.hws, img_window.width),
         )
-        preds = logits_a[:, : img_window.height, : img_window.width].softmax(axis=0)
+        preds = logits_a[:, : img_window.height, : img_window.width]
+
+        # Numpy softmax on axis 0
+        preds = np.exp(preds) / np.sum(np.exp(preds), axis=0)
 
         return preds, preds_win

kelp_o_matic/managers.py

@@ -2,13 +2,13 @@
 from pathlib import Path
 from typing import Union
 
+import numpy as np
 import rasterio
-import torch
 from rich import print
 from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
 
 from kelp_o_matic.geotiff_io import GeotiffReader, GeotiffWriter
-from kelp_o_matic.hann import TorchMemoryRegister, BartlettHannKernel
+from kelp_o_matic.hann import NumpyMemoryRegister, BartlettHannKernel
 from kelp_o_matic.models import _Model
 from kelp_o_matic.utils import all_same
 
@@ -54,9 +54,9 @@ def __init__(
             dtype="uint8",
             nodata=0,
         )
-        self.kernel = BartlettHannKernel(crop_size, self.model.device)
-        self.register = TorchMemoryRegister(
-            self.input_path, self.model.register_depth, crop_size, self.model.device
+        self.kernel = BartlettHannKernel(crop_size)
+        self.register = NumpyMemoryRegister(
+            self.input_path, self.model.register_depth, crop_size
         )
 
     def __call__(self):
@@ -74,17 +74,17 @@ def __call__(self):
                 if self.model.transform:
                     crop = self.model.transform(crop / self._max_value)
 
-                if torch.all(crop == 0):
+                if np.all(crop == 0):
                     logits = self.model.shortcut(self.reader.crop_size)
                 else:
                     # Zero pad to correct shape
-                    _, h, w = crop.shape
-                    crop = torch.nn.functional.pad(
-                        crop, (0, self.crop_size - w, 0, self.crop_size - h), value=0
-                    )
-                    logits = self.model(crop.unsqueeze(0))[0]
+                    c, h, w = crop.shape
+                    zeros = np.zeros((c, self.crop_size, self.crop_size), dtype=crop.dtype)
+                    zeros[:, :h, :w] = crop
+                    crop = zeros
+                    logits = self.model(np.expand_dims(crop,0))[0]
 
-                logits = self.kernel(
+                logits = self.kernel.forward(
                     logits,
                     top=self.reader.is_top_window(read_window),
                     bottom=self.reader.is_bottom_window(read_window),
@@ -198,10 +198,6 @@ def __call__(self):
         with self.progress:
             super().__call__()
 
-    def on_start(self):
-        device_emoji = ":rocket:" if self.model.device.type == "cuda" else ":snail:"
-        print(f"Running with [magenta]{self.model.device} {device_emoji}")
-
     def on_tile_write(self, index: int):
         self.progress.update(self.processing_task, completed=index)