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

change the size of input to remedy OOM issue.

Browse Source 
PiperOrigin-RevId: 317995769
Change-Id: I1358449e989a41c5621e6a4d56e603387be0490d
This commit is contained in:
Zhenyu Tan 2020-06-23 21:06:46 -07:00 committed by Geeta Chavan
parent e5a4b2534b
commit 83e4305ca4
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/python/keras/layers/preprocessing/image_preprocessing_distribution_test.py
View File

@ -40,9 +40,10 @@ class ImagePreprocessingDistributionTest(
preprocessing_test_utils.PreprocessingLayerTest): preprocessing_test_utils.PreprocessingLayerTest):
def test_distribution(self, distribution): def test_distribution(self, distribution):
np_images = np.random.random((1000, 32, 32, 3)).astype(np.float32) # TODO(b/159738418): large image input causes OOM in ubuntu multi gpu.
np_images = np.random.random((32, 32, 32, 3)).astype(np.float32)
image_dataset = dataset_ops.Dataset.from_tensor_slices(np_images).batch( image_dataset = dataset_ops.Dataset.from_tensor_slices(np_images).batch(
32, drop_remainder=True) 16, drop_remainder=True)
with distribution.scope(): with distribution.scope():
input_data = keras.Input(shape=(32, 32, 3), dtype=dtypes.float32) input_data = keras.Input(shape=(32, 32, 3), dtype=dtypes.float32)
@ -58,7 +59,7 @@ class ImagePreprocessingDistributionTest(
output = flatten_layer(preprocessed_image) output = flatten_layer(preprocessed_image)
cls_layer = keras.layers.Dense(units=1, activation="sigmoid") cls_layer = keras.layers.Dense(units=1, activation="sigmoid")
output = cls_layer(output) output = cls_layer(output)
model = keras.Model(inputs=input_data, outputs=preprocessed_image) model = keras.Model(inputs=input_data, outputs=output)
model.compile(loss="binary_crossentropy") model.compile(loss="binary_crossentropy")
_ = model.predict(image_dataset) _ = model.predict(image_dataset)