Merge pull request #136 from johnolafenwa/john/newer-gpu-support

Added Support for Newer GPUS

Author: johnolafenwa

Dockerfile.cpu

@@ -1,4 +1,4 @@
-FROM deepquestai/deepstack-base:cpu-412378893 as cpu 
+FROM deepquestai/deepstack-base:cpu-1647248158 as cpu 
 
 ENV SLEEP_TIME 0.01
 ENV CUDA_MODE False

Dockerfile.gpu

@@ -1,4 +1,4 @@
-FROM deepquestai/deepstack-base:gpu-412378893 as gpu
+FROM deepquestai/deepstack-base:gpu-1647248158 as gpu
 
 ENV SLEEP_TIME 0.01
 ENV CUDA_MODE True

intelligencelayer/shared/models/common.py

@@ -4,7 +4,7 @@
 import torch
 import torch.nn as nn
 from utils.general import non_max_suppression
-
+import torch.nn.functional as F
 
 def autopad(k, p=None):  # kernel, padding
     # Pad to 'same'
@@ -26,7 +26,7 @@ def __init__(
         super(Conv, self).__init__()
         self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
         self.bn = nn.BatchNorm2d(c2)
-        self.act = nn.Hardswish() if act else nn.Identity()
+        self.act = Hardswish() if act else nn.Identity()
 
     def forward(self, x):
         return self.act(self.bn(self.conv(x)))

intelligencelayer/shared/process.py

@@ -4,7 +4,7 @@
 import shutil
 import time
 from pathlib import Path
-
+import models
 import numpy as np
 import torch
 import torch.backends.cudnn as cudnn
@@ -23,7 +23,8 @@
     xyxy2xywh,
 )
 from utils.torch_utils import load_classifier, select_device, time_synchronized
-
+import torch.nn as nn
+from utils.activations import Hardswish
 
 class YOLODetector(object):
     def __init__(self, model_path: str, reso: int = 640, cuda: bool = False):
@@ -34,6 +35,13 @@ def __init__(self, model_path: str, reso: int = 640, cuda: bool = False):
         self.reso = (reso, reso)
         self.cuda = cuda
         self.model = attempt_load(model_path, map_location=self.device)
+
+        # Update model
+        for k, m in self.model.named_modules():
+            m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatability
+            if isinstance(m, models.common.Conv) and isinstance(m.act, nn.Hardswish):
+                m.act = Hardswish()  # assign activation
+
         self.names = (
             self.model.module.names
             if hasattr(self.model, "module")

intelligencelayer/shared/utils/activations.py

@@ -9,13 +9,29 @@ class Swish(nn.Module):  #
     def forward(x):
         return x * torch.sigmoid(x)
 
-
-class Hardswish(nn.Module):  # export-friendly version of nn.Hardswish()
-    @staticmethod
-    def forward(x):
-        # return x * F.hardsigmoid(x)  # for torchscript and CoreML
-        return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0  # for torchscript, CoreML and ONNX
-
+class Hardswish(nn.Module):
+    r"""Applies the hardswish function, element-wise, as described in the paper:
+    `Searching for MobileNetV3`_.
+    .. math::
+        \text{Hardswish}(x) = \begin{cases}
+            0 & \text{if~} x \le -3, \\
+            x & \text{if~} x \ge +3, \\
+            x \cdot (x + 3) /6 & \text{otherwise}
+        \end{cases}
+    Shape:
+        - Input: :math:`(N, *)` where `*` means, any number of additional
+          dimensions
+        - Output: :math:`(N, *)`, same shape as the input
+    Examples::
+        >>> m = nn.Hardswish()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    .. _`Searching for MobileNetV3`:
+        https://arxiv.org/abs/1905.02244
+    """
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.hardswish(input)
 
 class MemoryEfficientSwish(nn.Module):
     class F(torch.autograd.Function):