Ignore negative batch dimensions when calculating batch_size. (#173)

* Ignore negative batch dimensions and cast ratio as abs.

redis_consumer/grpc_clients.py

@@ -32,6 +32,7 @@
 
 import json
 import logging
+import math
 import time
 import timeit
 import six
@@ -348,12 +349,15 @@ def get_batch_size(self):
 
         ratio = 1
         for shape in input_shape:
-            rank = len(shape)
-            ratio *= (shape[rank - 3] / settings.TF_MIN_MODEL_SIZE) * \
-                     (shape[rank - 2] / settings.TF_MIN_MODEL_SIZE) * \
-                     (shape[rank - 1])
+            # reduce batch size by image dimensions, ignore batch
+            for size in shape[1:-1]:
+                if size > 0:  # ignore None dimensions (-1)
+                    ratio *= size / settings.TF_MIN_MODEL_SIZE
 
-        batch_size = int(settings.TF_MAX_BATCH_SIZE // ratio)
+            # reduce by channel
+            ratio *= shape[-1]
+
+        batch_size = int(settings.TF_MAX_BATCH_SIZE // math.fabs(ratio))
         return batch_size
 
     def predict(self, tiles, batch_size=None):

redis_consumer/grpc_clients_test.py

@@ -178,6 +178,13 @@ def test_get_batch_size(self, mocker):
             batch_size = wrapper.get_batch_size()
             assert batch_size == settings.TF_MAX_BATCH_SIZE * m * m
 
+        # what if the model has no defined size except the channel
+        channels = 3
+        shape = [(-1, -1, -1, -1, channels)]
+        mocker.patch.object(wrapper, 'input_shape', shape)
+        batch_size = wrapper.get_batch_size()
+        assert batch_size == settings.TF_MAX_BATCH_SIZE // channels
+
     def test_send_grpc(self):
         client = DummyPredictClient(1, 2, 3)
         metadata = self._get_metadata()