Colab: UnknownError: Failed to get convolution algorithm when using TF 2.5

Question

我在 Google Colab 上使用 tensorflow 2.5 时似乎遇到了问题。我假设 CUDA 版本 and/or CuDNN 版本之间存在一些不兼容性。我该如何修复它们？

我检查了colab使用的CUDA version。它是 11.2，tf2.5 应该没问题。这意味着问题出在 CuDNN 上，对吗？

重现代码：

!pip install tensorflow==2.5.0
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.datasets import cifar10

(x_train, y_train), (x_test, y_test) = cifar10.load_data()
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

def my_model():
    inputs = keras.Input(shape=(32, 32, 3))
    x = layers.Conv2D(32, 3)(inputs)
    x = layers.BatchNormalization()(x)
    x = keras.activations.relu(x)
    x = layers.MaxPooling2D()(x)
    x = layers.Conv2D(64, 3)(x)
    x = layers.BatchNormalization()(x)
    x = keras.activations.relu(x)
    x = layers.MaxPooling2D()(x)
    x = layers.Conv2D(128, 3)(x)
    x = layers.BatchNormalization()(x)
    x = keras.activations.relu(x)
    x = layers.Flatten()(x)
    x = layers.Dense(64, activation="relu")(x)
    outputs = layers.Dense(10)(x)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


model = my_model()
model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.Adam(learning_rate=3e-4),
    metrics=["accuracy"],
)

model.fit(x_train, y_train, batch_size=64, epochs=10, verbose=2)
model.evaluate(x_test, y_test, batch_size=64, verbose=2)

Error I get

我试过但我得到了同样的错误。

还建议我使用 tf.config.experimental.set_memory_growth(gpu, True) 但同样 - 这不起作用 - 我得到同样的错误。

我对使用 GPU 很感兴趣。我知道没有硬件加速一切正常。

Answer 1

在 this 文档中，Google 警告我们不要 install/downgrade 使用 !pip 命令的 TensorFlow 版本。
他们写道：

Colab builds TensorFlow from source to ensure compatibility with our fleet of accelerators. Versions of TensorFlow fetched from PyPI by pip may suffer from performance problems or may not work at all.

这意味着如果我们安装任何其他 TensorFlow，该版本可能与其提供的 GPU/TPU 配置不兼容。所以，只需使用 TensorFlow 2.6（这是最新版本），它与版本 2.5 非常相似。

Colab: UnknownError: Failed to get convolution algorithm when using TF 2.5

Colab: UnknownError: Failed to get convolution algorithm when using TF 2.5

tensorflow

cudnn

google-colaboratory