Merge pull request #153 from gnes-ai/feat-pool-encoder

feat(encoder): separate pooling as an indep. encoder

gnes/base/__init__.py

@@ -38,7 +38,7 @@ def register_all_class(cls2file_map: Dict, module_name: str):
             getattr(importlib.import_module('gnes.%s.%s' % (module_name, v)), k)
         except ImportError as ex:
             default_logger = set_logger('GNES')
-            default_logger.warning('fail to register %s, due to %s' % (k, ex))
+            default_logger.warning('fail to register %s, due to "%s", you will not be able to use this model' % (k, ex))
     load_contrib_module()
 
 
@@ -65,7 +65,8 @@ class TrainableType(type):
         'is_trained': False,
         'batch_size': None,
         'work_dir': os.environ.get('GNES_VOLUME', os.getcwd()),
-        'name': None
+        'name': None,
+        'on_gpu': False
     }
 
     def __new__(cls, *args, **kwargs):

gnes/encoder/__init__.py

@@ -21,10 +21,7 @@
     'BertEncoder': 'text.bert',
     'BertEncoderWithServer': 'text.bert',
     'BertEncoderServer': 'text.bert',
-    'ElmoEncoder': 'text.elmo',
     'FlairEncoder': 'text.flair',
-    'GPTEncoder': 'text.gpt',
-    'GPT2Encoder': 'text.gpt',
     'PCALocalEncoder': 'numeric.pca',
     'PQEncoder': 'numeric.pq',
     'TFPQEncoder': 'numeric.tf_pq',
@@ -42,7 +39,9 @@
     'CVAEEncoder': 'image.cvae',
     'IncepMixtureEncoder': 'video.incep_mixture',
     'VladEncoder': 'numeric.vlad',
-    'MfccEncoder': 'audio.mfcc'
+    'MfccEncoder': 'audio.mfcc',
+    'PoolingEncoder': 'numeric.pooling',
+    'PyTorchTransformers': 'text.transformer'
 }
 
 register_all_class(_cls2file_map, 'encoder')

gnes/encoder/base.py

@@ -14,7 +14,7 @@
 #  limitations under the License.
 
 
-from typing import List, Any
+from typing import List, Any, Tuple, Union
 
 import numpy as np
 
@@ -44,13 +44,13 @@ def encode(self, data: List['np.ndarray'], *args, **kwargs) -> np.ndarray:
 
 class BaseTextEncoder(BaseEncoder):
 
-    def encode(self, text: List[str], *args, **kwargs) -> np.ndarray:
+    def encode(self, text: List[str], *args, **kwargs) -> Union[Tuple, np.ndarray]:
         pass
 
 
 class BaseNumericEncoder(BaseEncoder):
 
-    def encode(self, text: np.ndarray, *args, **kwargs) -> np.ndarray:
+    def encode(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
         pass
 
 

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/numeric/pooling.py

@@ -0,0 +1,100 @@
+import os
+from typing import Tuple
+
+import numpy as np
+
+from ..base import BaseNumericEncoder
+from ...helper import as_numpy_array
+
+
+class PoolingEncoder(BaseNumericEncoder):
+    def __init__(self, pooling_strategy: str = 'REDUCE_MEAN',
+                 backend: str = 'numpy',
+                 *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        valid_poolings = {'REDUCE_MEAN', 'REDUCE_MAX', 'REDUCE_MEAN_MAX'}
+        valid_backends = {'tensorflow', 'numpy', 'pytorch', 'torch'}
+
+        if pooling_strategy not in valid_poolings:
+            raise ValueError('"pooling_strategy" must be one of %s' % valid_poolings)
+        if backend not in valid_backends:
+            raise ValueError('"backend" must be one of %s' % valid_backends)
+        self.pooling_strategy = pooling_strategy
+        self.backend = backend
+
+    def post_init(self):
+        if self.backend in {'pytorch', 'torch'}:
+            import torch
+            self.torch = torch
+        elif self.backend == 'tensorflow':
+            os.environ['CUDA_VISIBLE_DEVICES'] = '0' if self.on_gpu else '-1'
+            import tensorflow as tf
+            self._tf_graph = tf.Graph()
+            config = tf.ConfigProto(device_count={'GPU': 1 if self.on_gpu else 0})
+            config.gpu_options.allow_growth = True
+            config.log_device_placement = False
+            self._sess = tf.Session(graph=self._tf_graph, config=config)
+            self.tf = tf
+
+    def mul_mask(self, x, m):
+        if self.backend in {'pytorch', 'torch'}:
+            return self.torch.mul(x, m.unsqueeze(2))
+        elif self.backend == 'tensorflow':
+            with self._tf_graph.as_default():
+                return x * self.tf.expand_dims(m, axis=-1)
+        elif self.backend == 'numpy':
+            return x * np.expand_dims(m, axis=-1)
+
+    def minus_mask(self, x, m, offset: int = 1e30):
+        if self.backend in {'pytorch', 'torch'}:
+            return x - (1.0 - m).unsqueeze(2) * offset
+        elif self.backend == 'tensorflow':
+            with self._tf_graph.as_default():
+                return x - self.tf.expand_dims(1.0 - m, axis=-1) * offset
+        elif self.backend == 'numpy':
+            return x - np.expand_dims(1.0 - m, axis=-1) * offset
+
+    def masked_reduce_mean(self, x, m, jitter: float = 1e-10):
+        if self.backend in {'pytorch', 'torch'}:
+            return self.torch.div(self.torch.sum(self.mul_mask(x, m), dim=1),
+                                  self.torch.sum(m.unsqueeze(2), dim=1) + jitter)
+        elif self.backend == 'tensorflow':
+            with self._tf_graph.as_default():
+                return self.tf.reduce_sum(self.mul_mask(x, m), axis=1) / (
+                        self.tf.reduce_sum(m, axis=1, keepdims=True) + jitter)
+        elif self.backend == 'numpy':
+            return np.sum(self.mul_mask(x, m), axis=1) / (np.sum(m, axis=1, keepdims=True) + jitter)
+
+    def masked_reduce_max(self, x, m):
+        if self.backend in {'pytorch', 'torch'}:
+            return self.torch.max(self.minus_mask(x, m), 1)[0]
+        elif self.backend == 'tensorflow':
+            with self._tf_graph.as_default():
+                return self.tf.reduce_max(self.minus_mask(x, m), axis=1)
+        elif self.backend == 'numpy':
+            return np.max(self.minus_mask(x, m), axis=1)
+
+    @as_numpy_array
+    def encode(self, data: Tuple, *args, **kwargs):
+        seq_tensor, mask_tensor = data
+
+        if self.pooling_strategy == 'REDUCE_MEAN':
+            r = self.masked_reduce_mean(seq_tensor, mask_tensor)
+        elif self.pooling_strategy == 'REDUCE_MAX':
+            r = self.masked_reduce_max(seq_tensor, mask_tensor)
+        elif self.pooling_strategy == 'REDUCE_MEAN_MAX':
+            if self.backend in {'pytorch', 'torch'}:
+                r = self.torch.cat((self.masked_reduce_mean(seq_tensor, mask_tensor),
+                                    self.masked_reduce_max(seq_tensor, mask_tensor)), dim=1)
+            elif self.backend == 'tensorflow':
+                with self._tf_graph.as_default():
+                    r = self.tf.concat([self.masked_reduce_mean(seq_tensor, mask_tensor),
+                                        self.masked_reduce_max(seq_tensor, mask_tensor)], axis=1)
+            elif self.backend == 'numpy':
+                r = np.concatenate([self.masked_reduce_mean(seq_tensor, mask_tensor),
+                                    self.masked_reduce_max(seq_tensor, mask_tensor)], axis=1)
+
+        if self.backend == 'tensorflow':
+            r = self._sess.run(r)
+        return r

gnes/encoder/text/flair.py

@@ -19,34 +19,30 @@
 import numpy as np
 
 from ..base import BaseTextEncoder
-from ...helper import batching, pooling_np
+from ...helper import batching, as_numpy_array
 
 
 class FlairEncoder(BaseTextEncoder):
     is_trained = True
 
-    def __init__(self, model_name: str = 'multi-forward-fast',
-                 pooling_strategy: str = 'REDUCE_MEAN', *args, **kwargs):
+    def __init__(self, pooling_strategy: str = 'mean', *args, **kwargs):
         super().__init__(*args, **kwargs)
-
-        self.model_name = model_name
         self.pooling_strategy = pooling_strategy
 
     def post_init(self):
-        from flair.embeddings import FlairEmbeddings
-        self._flair = FlairEmbeddings(self.model_name)
+        from flair.embeddings import DocumentPoolEmbeddings, WordEmbeddings, FlairEmbeddings
+        self._flair = DocumentPoolEmbeddings(
+            [WordEmbeddings('glove'),
+             FlairEmbeddings('news-forward'),
+             FlairEmbeddings('news-backward')],
+            pooling=self.pooling_strategy)
 
     @batching
+    @as_numpy_array
     def encode(self, text: List[str], *args, **kwargs) -> np.ndarray:
         from flair.data import Sentence
+        import torch
         # tokenize text
-        batch_tokens = [Sentence(sent) for sent in text]
-
-        flair_encodes = self._flair.embed(batch_tokens)
-
-        pooled_data = []
-        for sentence in flair_encodes:
-            _layer_data = np.stack([s.embedding.numpy() for s in sentence])
-            _pooled = pooling_np(_layer_data, self.pooling_strategy)
-            pooled_data.append(_pooled)
-        return np.array(pooled_data, dtype=np.float32)
+        batch_tokens = [Sentence(v) for v in text]
+        self._flair.embed(batch_tokens)
+        return torch.stack([v.embedding for v in batch_tokens]).detach()