diff --git a/oneflow/python/test/ops/test_batch_normalization.py b/oneflow/python/test/ops/test_batch_normalization.py index 28703d424c4b919b90a16d69fa11bf1a795a0f62..1d6eb24ef7906d8724deef67dc4b18d411b0fc95 100644 --- a/oneflow/python/test/ops/test_batch_normalization.py +++ b/oneflow/python/test/ops/test_batch_normalization.py @@ -24,6 +24,10 @@ from test_util import Args, GenArgDict, type_name_to_flow_type, type_name_to_np_ import oneflow.typing as oft import unittest +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def TODO_test_train(test_case): flow.config.enable_debug_mode(True) diff --git a/oneflow/python/test/ops/test_broadcast_like.py b/oneflow/python/test/ops/test_broadcast_like.py index 55df6ea057ad7bda6630a0020c508a7580b1a7a8..edc93cf46d69da7944d0541d6648729e89868881 100644 --- a/oneflow/python/test/ops/test_broadcast_like.py +++ b/oneflow/python/test/ops/test_broadcast_like.py @@ -23,6 +23,10 @@ import tensorflow as tf from test_util import GenArgList import oneflow.typing as tp +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def compare_broadcast_like_with_tf( device_type, input_shape, like_shape, broadcast_axes, rtol=1e-5, atol=1e-5 diff --git a/oneflow/python/test/ops/test_cast.py b/oneflow/python/test/ops/test_cast.py index 2226b3dc85e0d035600083c1c1682e4e809cb1a2..3875396b87ee5a781be45976ebf8378db80f61ab 100644 --- a/oneflow/python/test/ops/test_cast.py +++ b/oneflow/python/test/ops/test_cast.py @@ -24,6 +24,10 @@ import test_global_storage from test_util import GenArgList, type_name_to_flow_type, type_name_to_np_type import oneflow.typing as oft +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def cast_forward_compare_with_tensorflow(test_cast, device_type, input_shape, dtype): assert device_type in ["gpu", "cpu"] diff --git a/oneflow/python/test/ops/test_deconv2d.py b/oneflow/python/test/ops/test_deconv2d.py index 56cfb637f1ec6da532867f2b6d392a7f9b081833..0d403ee5a4fdbc8b75ced7ac316f680a1097d90e 100644 --- a/oneflow/python/test/ops/test_deconv2d.py +++ b/oneflow/python/test/ops/test_deconv2d.py @@ -23,6 +23,10 @@ import tensorflow as tf import test_global_storage from test_util import GenArgList +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def compare_with_tensorflow(device_type, params_case, dilations, data_format): input_shape, output_shape, padding, strides, kernel_size = params_case diff --git a/oneflow/python/test/ops/test_fuse_cast_scale.py b/oneflow/python/test/ops/test_fuse_cast_scale.py index 5a2e530faec39d9322e2457b86377e7378a3cbd6..7c5249e367aacbf3404637fe349a05ace681c3c4 100644 --- a/oneflow/python/test/ops/test_fuse_cast_scale.py +++ b/oneflow/python/test/ops/test_fuse_cast_scale.py @@ -24,6 +24,10 @@ import test_global_storage from test_util import GenArgList, type_name_to_flow_type, type_name_to_np_type import oneflow.typing as oft +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def fused_cast_scale(x, scalar, name): return ( diff --git a/oneflow/python/test/ops/test_reduce_ops.py b/oneflow/python/test/ops/test_reduce_ops.py index cbcfef2c307a7910543d10716ccf278c66fcafb9..1ef5f83672ffa579c4bdb2ace9ff0baaf23f9684 100644 --- a/oneflow/python/test/ops/test_reduce_ops.py +++ b/oneflow/python/test/ops/test_reduce_ops.py @@ -25,6 +25,10 @@ from test_util import GenArgList import oneflow.typing as oft import test_global_storage +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def compare_reduce_any_with_tensorflow( device_type, input_shape, axis, keepdims, rtol=1e-5, atol=1e-5 diff --git a/oneflow/python/test/ops/test_reduce_opsV2.py b/oneflow/python/test/ops/test_reduce_opsV2.py index f5ef7efd7399d990e825d5a865936ece87eda903..24ef1d36b43efb0fca8488b5f7a0e7861b5a404c 100644 --- a/oneflow/python/test/ops/test_reduce_opsV2.py +++ b/oneflow/python/test/ops/test_reduce_opsV2.py @@ -25,6 +25,10 @@ import test_global_storage from test_util import GenArgList import oneflow.typing as oft +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def compare_reduce_sum_with_tensorflow( device_type, input_shape, axis, keepdims, rtol=1e-5, atol=1e-5 diff --git a/oneflow/python/test/ops/test_reshapeV2.py b/oneflow/python/test/ops/test_reshapeV2.py index 78dc528a64a27b15b00fa31e126567fa7b546ac2..3f9ab18fa3d195e53d74a8d43302177c71a9ad1d 100644 --- a/oneflow/python/test/ops/test_reshapeV2.py +++ b/oneflow/python/test/ops/test_reshapeV2.py @@ -23,6 +23,10 @@ import tensorflow as tf import test_global_storage from test_util import GenArgList +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def compare_with_tensorflow(device_type, input_shape, shape): assert device_type in ["gpu", "cpu"] diff --git a/oneflow/python/test/ops/test_summary.py b/oneflow/python/test/ops/test_summary.py index abfbb915308413ee0f84e72659ac02d44f959708..55643f1b908cfe76cf1b6a00ea3984338f1c6b23 100644 --- a/oneflow/python/test/ops/test_summary.py +++ b/oneflow/python/test/ops/test_summary.py @@ -15,7 +15,6 @@ limitations under the License. """ import os import numpy as np -import tensorflow as tf import oneflow as flow from collections import OrderedDict import cv2 diff --git a/oneflow/python/test/ops/test_transpose.py b/oneflow/python/test/ops/test_transpose.py index 23c20a474047360912602094a7f2503662ed999b..a6f80145b0e4e63eb5bcc7eec7fe6463012804b4 100644 --- a/oneflow/python/test/ops/test_transpose.py +++ b/oneflow/python/test/ops/test_transpose.py @@ -24,6 +24,10 @@ import tensorflow as tf import test_global_storage from test_util import GenArgList +gpus = tf.config.experimental.list_physical_devices("GPU") +for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + def compare_with_tensorflow(device_type, input_shape, perm): assert device_type in ["gpu", "cpu"] diff --git a/oneflow/python/test/ops/test_util.py b/oneflow/python/test/ops/test_util.py index 67860396e22caebeabfac14148048d3c6c6e6cfa..47038883c4b143e3b5d76935a61b1cecabe90a4c 100644 --- a/oneflow/python/test/ops/test_util.py +++ b/oneflow/python/test/ops/test_util.py @@ -95,6 +95,10 @@ def RunOneflowOp(device_type, flow_op, x, flow_args): def RunTensorFlowOp(tf_op, x, tf_args): import tensorflow as tf + gpus = tf.config.experimental.list_physical_devices("GPU") + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + with tf.GradientTape(persistent=True) as tape: x = tf.Variable(x) y = tf_op(x, *tf_args) diff --git a/oneflow/python/test/ops/test_watch_diff.py b/oneflow/python/test/ops/test_watch_diff.py index c93bc13893d21a457ba87713938fbde2e43e2eb6..8054d7ea9060b303f617fa342fb00563e5048cfe 100644 --- a/oneflow/python/test/ops/test_watch_diff.py +++ b/oneflow/python/test/ops/test_watch_diff.py @@ -19,7 +19,6 @@ from collections import OrderedDict import numpy as np import oneflow as flow -import tensorflow as tf import test_global_storage from test_util import GenArgList, type_name_to_flow_type, type_name_to_np_type