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

Merge pull request #40803 from geetachavan1/cherrypicks_7XBWE

Browse Source 
[CherryPick 2.3] change the size of input to remedy OOM issue.
This commit is contained in:
Goldie Gadde 2020-06-25 12:01:51 -07:00 committed by GitHub
commit 890eae3e88
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
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):
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(
32, drop_remainder=True)
16, drop_remainder=True)
with distribution.scope():
input_data = keras.Input(shape=(32, 32, 3), dtype=dtypes.float32)
@ -58,7 +59,7 @@ class ImagePreprocessingDistributionTest(
output = flatten_layer(preprocessed_image)
cls_layer = keras.layers.Dense(units=1, activation="sigmoid")
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.predict(image_dataset)