not handling batch shape correctly led to #5

henrysky · May 14, 2018 · 4be8d5a · 4be8d5a
1 parent 64b17b0
commit 4be8d5a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/astroNN/models/BayesianCNNBase.py b/astroNN/models/BayesianCNNBase.py
@@ -354,30 +354,33 @@ def test(self, input_data, inputs_err=None):
 
         # for number of training data smaller than batch_size
         if total_test_num < self.batch_size:
-            self.batch_size = total_test_num
+            batch_size = total_test_num
+        else:
+            batch_size = self.batch_size
 
         # Due to the nature of how generator works, no overlapped prediction
-        data_gen_shape = (total_test_num // self.batch_size) * self.batch_size
+        data_gen_shape = (total_test_num // batch_size) * batch_size
         remainder_shape = total_test_num - data_gen_shape  # Remainder from generator
 
         start_time = time.time()
         print("Starting Dropout Variational Inference")
 
         # Data Generator for prediction
-        prediction_generator = BayesianCNNPredDataGenerator(self.batch_size).generate(input_array[:data_gen_shape],
-                                                                                      inputs_err[:data_gen_shape])
+        prediction_generator = BayesianCNNPredDataGenerator(batch_size).generate(input_array[:data_gen_shape],
+                                                                                 inputs_err[:data_gen_shape])
 
         new = FastMCInference(self.mc_num)(self.keras_model_predict)
 
-        result = np.asarray(new.predict_generator(prediction_generator, steps=data_gen_shape // self.batch_size))
+        result = np.asarray(new.predict_generator(prediction_generator, steps=data_gen_shape // batch_size))
 
         if remainder_shape != 0:  # deal with remainder
             remainder_generator = BayesianCNNPredDataGenerator(remainder_shape).generate(input_array[data_gen_shape:],
                                                                                           inputs_err[data_gen_shape:])
             remainder_result = np.asarray(new.predict_generator(remainder_generator, steps=1))
             result = np.concatenate((result, remainder_result))
 
-        if result.ndim < 3:  # in case only 1 test data point, in such case we need to add a dimension
+        # in case only 1 test data point, in such case we need to add a dimension
+        if result.ndim < 3 and batch_size == 1:
             result = np.expand_dims(result, axis=0)
 
         half_first_dim = result.shape[1] // 2  # result.shape[1] is guarantee an even number, otherwise sth is wrong

diff --git a/astroNN/nn/layers.py b/astroNN/nn/layers.py
@@ -461,7 +461,7 @@ def call(self, inputs, training=None):
         """
         # need to stack because keras can only handle one output
         mean, var = tf.nn.moments(inputs, axes=1)
-        return tf.squeeze(tf.stack([[mean], [var]], axis=-1))
+        return tf.stack((tf.squeeze([mean]), tf.squeeze([var])), axis=-1)
 
 
 class FastMCRepeat(Layer):