refactor(base): add on_gpu to replace use_cuda

gnes/encoder/image/inception.py

@@ -27,13 +27,11 @@ class TFInceptionEncoder(BaseImageEncoder):
 
     def __init__(self, model_dir: str,
                  select_layer: str = 'PreLogitsFlatten',
-                 use_cuda: bool = False,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         self.model_dir = model_dir
         self.select_layer = select_layer
-        self._use_cuda = use_cuda
         self.inception_size_x = 299
         self.inception_size_y = 299
 
@@ -57,7 +55,7 @@ def post_init(self):
                                                             dropout_keep_prob=1.0)
 
             config = tf.ConfigProto(log_device_placement=False)
-            if self._use_cuda:
+            if self.on_gpu:
                 config.gpu_options.allow_growth = True
             self.sess = tf.Session(config=config)
             self.saver = tf.train.Saver()

gnes/encoder/image/onnx.py

@@ -26,13 +26,11 @@ class BaseONNXImageEncoder(BaseImageEncoder):
 
     def __init__(self, model_name: str,
                  model_dir: str,
-                 use_cuda: bool = False,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         self.model_dir = model_dir
         self.model_name = model_name
-        self._use_cuda = use_cuda
 
     def post_init(self):
         import onnxruntime as ort

gnes/encoder/image/torchvision.py

@@ -19,7 +19,7 @@
 import numpy as np
 
 from ..base import BaseImageEncoder
-from ...helper import batching
+from ...helper import batching, as_numpy_array
 
 
 class TorchvisionEncoder(BaseImageEncoder):
@@ -28,14 +28,12 @@ class TorchvisionEncoder(BaseImageEncoder):
     def __init__(self, model_name: str,
                  layers: List[str],
                  model_dir: str,
-                 use_cuda: bool = False,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         self.model_dir = model_dir
         self.model_name = model_name
         self.layers = layers
-        self._use_cuda = use_cuda
 
     def post_init(self):
         import torch
@@ -69,7 +67,7 @@ def forward(self, x):
         os.environ['TORCH_HOME'] = self.model_dir
         self._model = _Model(self.model_name, self.layers)
         self._model = self._model.eval()
-        if self._use_cuda:
+        if self.on_gpu:
             # self._model.cuda()
             self._device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
             self._model = self._model.to(self._device)
@@ -94,6 +92,7 @@ def _padding(img: List['np.ndarray']):
             max_lenth = -1
 
         @batching(chunk_dim=max_lenth)
+        @as_numpy_array
         def _encode(_, img: List['np.ndarray']):
             import copy
 
@@ -106,14 +105,11 @@ def _encode(_, img: List['np.ndarray']):
                 img_for_torch = np.array(img, dtype=np.float32).transpose(0, 3, 1, 2)
 
             img_tensor = torch.from_numpy(img_for_torch)
-            if self._use_cuda:
+            if self.on_gpu:
                 img_tensor = img_tensor.cuda()
 
             encodes = self._model(img_tensor)
 
-            output = np.array(encodes.data.cpu().numpy(), dtype=np.float32)
-            return output
+            return encodes.data.cpu()
 
-        output = _encode(self, img)
-
-        return output
+        return _encode(self, img)

gnes/encoder/video/incep_mixture.py

@@ -28,7 +28,6 @@ class IncepMixtureEncoder(BaseVideoEncoder):
     def __init__(self, model_dir_inception: str,
                  model_dir_mixture: str,
                  select_layer: str = 'PreLogitsFlatten',
-                 use_cuda: bool = False,
                  feature_size: int = 300,
                  vocab_size: int = 28,
                  cluster_size: int = 256,
@@ -42,7 +41,6 @@ def __init__(self, model_dir_inception: str,
         self.model_dir_inception = model_dir_inception
         self.model_dir_mixture = model_dir_mixture
         self.select_layer = select_layer
-        self.use_cuda = use_cuda
         self.cluster_size = cluster_size
         self.feature_size = feature_size
         self.vocab_size = vocab_size
@@ -76,7 +74,7 @@ def post_init(self):
                                                             dropout_keep_prob=1.0)
 
             config = tf.ConfigProto(log_device_placement=False)
-            if self.use_cuda:
+            if self.on_gpu:
                 config.gpu_options.allow_growth = True
             self.sess = tf.Session(config=config)
             self.saver = tf.train.Saver()
@@ -85,7 +83,7 @@ def post_init(self):
         g2 = tf.Graph()
         with g2.as_default():
             config = tf.ConfigProto(log_device_placement=False)
-            if self.use_cuda:
+            if self.on_gpu:
                 config.gpu_options.allow_growth = True
             self.sess2 = tf.Session(config=config)
             self.mix_model = NetFV(feature_size=self.feature_size,

gnes/helper.py

@@ -422,8 +422,6 @@ def arg_wrapper(self, data, label=None, *args, **kwargs):
                             reduced_result[col] = reduced_result[col].reshape(
                                 (-1, chunk_dim, reduced_result[col].shape[1]))
                     final_result = tuple(reduced_result)
-                else:
-                    raise TypeError('dont know how to reduce %s' % type(final_result[0]))
 
             if len(final_result):
                 return final_result

gnes/preprocessor/image/segmentation.py

@@ -29,12 +29,10 @@ class SegmentPreprocessor(SizedPreprocessor):
 
     def __init__(self, model_name: str,
                  model_dir: str,
-                 _use_cuda: bool = False,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.model_name = model_name
         self.model_dir = model_dir
-        self._use_cuda = _use_cuda
 
     def post_init(self):
         import torch
@@ -43,7 +41,7 @@ def post_init(self):
         os.environ['TORCH_HOME'] = self.model_dir
         self._model = getattr(models.detection, self.model_name)(pretrained=True)
         self._model = self._model.eval()
-        if self._use_cuda:
+        if self.on_gpu:
             # self._model.cuda()
             self._device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
             self._model = self._model.to(self._device)
@@ -54,7 +52,7 @@ def apply(self, doc: 'gnes_pb2.Document'):
             original_image = Image.open(io.BytesIO(doc.raw_bytes))
             all_subareas, index = get_all_subarea(original_image)
             image_tensor = torch_transform(original_image)
-            if self._use_cuda:
+            if self.on_gpu:
                 image_tensor = image_tensor.cuda()
 
             seg_output = self._model([image_tensor])

tests/test_pytorch_transformers_encoder.py

@@ -22,7 +22,7 @@ def setUp(self):
     def test_encoding(self):
         vec = self.tt_encoder.encode(self.test_str)
         self.assertEqual(vec.shape[0], len(self.test_str))
-        self.assertEqual(vec.shape[2], 768)
+        self.assertEqual(vec.shape[1], 768)
 
     def test_dump_load(self):
         self.tt_encoder.dump(self.dump_path)
@@ -31,7 +31,7 @@ def test_dump_load(self):
 
         vec = tt_encoder2.encode(self.test_str)
         self.assertEqual(vec.shape[0], len(self.test_str))
-        self.assertEqual(vec.shape[2], 768)
+        self.assertEqual(vec.shape[1], 768)
 
     def tearDown(self):
         if os.path.exists(self.dump_path):

tests/yaml/base-segmentation-image-prep.yml

@@ -4,7 +4,6 @@ components:
     parameters:
       model_dir: ${FASTERRCNN_MODEL}
       model_name: fasterrcnn_resnet50_fpn
-      _use_cuda: false
     gnes_config:
       name: fasterRCNN
       is_trained: true

tutorials/component-yaml-spec.md

@@ -130,7 +130,6 @@ def __init__(self, model_name: str,
                  layers: List[str],
                  model_dir: str,
                  batch_size: int = 64,
-                 use_cuda: bool = False,
                  *args, **kwargs):
   # do model init...
   # ...

yaml-example/component/img_preprocessor_fasterRCNN.yml

@@ -3,6 +3,5 @@ parameters:
   model_dir: /ext_data/image_preprocessor
   model_name: fasterrcnn_resnet50_fpn
   target_img_size: 224
-  _use_cuda: false
 gnes_config:
   is_trained: true