session crashes (CPU->GPU Memcpy failed) when running BatchNormalization

[creakseek]

Issue Type

Bug

Source

binary

Tensorflow Version

tfnightly

Custom Code

Yes

OS Platform and Distribution

Linux Ubuntu 20.04

Mobile device

No response

Python version

3.9

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

When I run the code snippet below (relu and BatchNormalization) for the first time, I encounter InternalError. Then I try to run it again and session crashes.

Standalone code to reproduce the issue

1. Run for a single time: throws InternalError

import tensorflow as tf
import numpy as np
print(tf.__version__)
input_data = np.random.rand(1, 3, 3, 1).astype(np.float32)
output_data = tf.keras.activations.relu(
    tf.keras.layers.BatchNormalization(axis=-1)(input_data))

Log output:

2.11.0-dev20220919
InternalError: Exception encountered when calling layer 'batch_normalization' (type BatchNormalization).

{{function_node __wrapped__FusedBatchNormV3_device_/job:localhost/replica:0/task:0/device:GPU:0}} cuDNN launch failure : input shape ([1,3,3,1]) [Op:FusedBatchNormV3]

Call arguments received by layer 'batch_normalization' (type BatchNormalization):
  • inputs=tf.Tensor(shape=(1, 3, 3, 1), dtype=float32)
  • training=None

2. Run for a second time: crash

import tensorflow as tf
import numpy as np
try:
  input_data = np.random.rand(1, 3, 3, 1).astype(np.float32)
  output_data = tf.keras.activations.relu(
    tf.keras.layers.BatchNormalization(axis=-1)(input_data))
except:
  pass

input_data = np.random.rand(1, 3, 3, 1).astype(np.float32)
output_data = tf.keras.activations.relu(
    tf.keras.layers.BatchNormalization(axis=-1)(input_data)) # crash here

Relevant logs:

F tensorflow/core/common_runtime/gpu/gpu_util.cc:386] CPU->GPU Memcpy failed

[tilakrayal]

@creakseek,

When running on GPU (with cuDNN version 7.6.3 or later), tf.nn.depthwise_conv2d backprop to filter (and therefore also tf.keras.layers.DepthwiseConv2D) now operate deterministically (and tf.errors.UnimplementedError is no longer thrown) when op-determinism has been enabled via tf.config.experimental.enable_op_determinism.

Could you please try to run the below command before executing the code on tensorflow v2.10. I tried and was able to execute the code. Kindly find the gist of it here.

!apt install --allow-change-held-packages libcudnn8=8.1.0.77-1+cuda11.2

Thank you!

[google-ml-butler]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[creakseek]

Thank you @tilakrayal ! That solves the problem.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No