TensorFlow 2.9.1静默升级致Google Colab GPU运行时Notebook故障

近三天内，Google Colab将TensorFlow从2.8.x版本静默升级至2.9.1，导致所有研究Notebook（包括下文的MNIST极简示例）出现故障。查阅TensorFlow发布说明后，未发现所用Keras或TensorFlow相关组件存在变更。进一步排查确认：该错误仅在Colab配置GPU运行时触发，使用CPU或TPU运行时则完全正常。

故障复现代码 #

import tensorflow as tf
import keras

# the data, split between train and test sets
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
x_test  = x_test.reshape ( x_test.shape[0], 28, 28, 1)
input_shape = (28, 28, 1)

x_train  = x_train.astype('float32')
x_test   = x_test.astype('float32')
x_train /= 255
x_test  /= 255

# convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, 10)
y_test  = keras.utils.to_categorical(y_test , 10)

model = keras.models.Sequential()
model.add(keras.layers.Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape))
model.add(keras.layers.Conv2D(64, (3, 3), activation='relu'))
model.add(keras.layers.MaxPooling2D(pool_size=(2, 2)))
model.add(keras.layers.Dropout(0.25))
model.add(keras.layers.Flatten())
model.add(keras.layers.Dense(128, activation='relu'))
model.add(keras.layers.Dropout(0.5))
model.add(keras.layers.Dense(10, activation='softmax'))

model.compile(loss=keras.losses.categorical_crossentropy, optimizer=keras.optimizers.Adam(),
              metrics=['accuracy'])

model.fit(x_train, y_train, batch_size=100, epochs=1, verbose=1, validation_data=(x_test, y_test))
score = model.evaluate(x_test, y_test, verbose=0)
print('Test loss:', score[0], 'Test accuracy:', score[1])

GPU运行时错误日志 #

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11490434/11490434 [==============================] - 0s 0us/step
---------------------------------------------------------------------------
UnimplementedError                        Traceback (most recent call last)
[<ipython-input-1-05f207168698>] in <module>
     31               metrics=['accuracy'])
     32 
---> 33 model.fit(x_train, y_train, batch_size=100, epochs=1, verbose=1, validation_data=(x_test, y_test))
     34 score = model.evaluate(x_test, y_test, verbose=0)
     35 print('Test loss:', score[0], 'Test accuracy:', score[1])

1 frames
[/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py] in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     53     ctx.ensure_initialized()
     54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:
     57     if name is not None:

UnimplementedError: Graph execution error:

    [...]

Node: 'sequential/conv2d/Conv2D'
DNN library is not found.
     [[{{node sequential/conv2d/Conv2D}}]] [Op:__inference_train_function_865]

临时解决方案 #

将TensorFlow降级至2.8.2版本后，故障完全消失，所有Colab Notebook恢复正常运行。以下是GPU运行时环境下的临时修复代码，注意每次运行会增加约86秒的耗时：

!pip install tensorflow==2.8.2
import tensorflow as tf
print(tf.__version__)

内容的提问来源于stack exchange，提问作者vwvan