experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update to use tf.shape to get the shape of the tensor, from review comment

Browse Source 
Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
This commit is contained in:
Yong Tang 2020-05-04 17:06:51 +00:00
parent d5598f1609
commit 1f44a4cd05
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tensorflow/python/keras/layers/normalization.py
View File

@ -736,10 +736,13 @@ class BatchNormalizationBase(Layer):
if self.virtual_batch_size is not None:
# Virtual batches (aka ghost batches) can be simulated by reshaping the
# Tensor and reusing the existing batch norm implementation
original_shape = [
d if d is not None else -1 for d in inputs.shape.as_list()]
original_shape = [-1] + original_shape[1:]
expanded_shape = [self.virtual_batch_size, -1] + original_shape[1:]
original_shape = array_ops.shape(inputs)
original_shape = array_ops.concat([
constant_op.constant([-1]),
original_shape[1:]], axis=0)
expanded_shape = array_ops.concat([
constant_op.constant([self.virtual_batch_size, -1]),
original_shape[1:]], axis=0)
# Will cause errors if virtual_batch_size does not divide the batch size
inputs = array_ops.reshape(inputs, expanded_shape)