Tensorflow：GPU上矩阵乘法（NaN）的错误结果

Posted 2023-02-16

【中文标题】Tensorflow：GPU上矩阵乘法（NaN）的错误结果

【英文标题】：Tensorflow: incorrect result of matrix multiplication (NaN) on GPU

【发布时间】：2019-10-29 14:59:15

【问题描述】：

tf.matmul (u_hat = tf.matmul(W_tiled, u_tiled)) 在 CPU 和 GPU 上返回不同的结果。第二批后 tensorflow/numpy matmul 操作的平均值为 3.6066954e+17/2.7731653e-06。最后，当代码在 GPU 上运行时，矩阵乘积开始包含 NaN 值。

复制于：

赢 10 | GTX 1080 | CUDA 工具包 10.0 | Python 3.7 | GPU 驱动程序 431.86 | TensorFlow 2.0

赢 10 | GTX 1080 | CUDA 工具包 10.0 | Python 3.7 | GPU 驱动程序 431.86 | TensorFlow 1.15

赢 10 | GTX 1080 | CUDA 工具包 10.0 | Python 3.7 | GPU 驱动程序 441.08 | TensorFlow 2.0

赢 10 | GTX 1080 | CUDA 工具包 10.0 | Python 3.7 | GPU 驱动程序 441.08 | TensorFlow 1.15

from tensorflow.keras import backend as K
from keras.layers import Layer
from tensorflow.keras import utils
from tensorflow.keras.datasets import mnist
import numpy as np
from tensorflow.keras.layers import *
import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import tensorflow as tf


def squash(s, axis=-1, epsilon=1e-7, name=None):
    squared_norm = tf.reduce_sum(tf.square(s), axis=axis, keepdims=True)
    safe_norm = tf.sqrt(squared_norm + epsilon)
    squash_factor = squared_norm / (1. + squared_norm)
    unit_vector = s / safe_norm
    return squash_factor * unit_vector


def softmax(x, axis=-1):
    ex = K.exp(x - K.max(x, axis=axis, keepdims=True))
    return ex / K.sum(ex, axis=axis, keepdims=True)


def safe_norm(s, axis=-1, epsilon=1e-7, keep_dims=False, name=None):
    squared_norm = tf.reduce_sum(tf.square(s), axis=axis, keepdims=keep_dims)
    return tf.sqrt(squared_norm + epsilon)


class Custom_layer(Layer):
    def __init__(self, **kwargs):
        super(Custom_layer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.kernel = self.add_weight(
            shape=(1, 1152, 10, 16, 8),
            initializer=tf.keras.initializers.RandomNormal(0.0, stddev=0.01),
            trainable=True)

    def call(self, inputs):
        reshaped = tf.reshape(inputs, [-1, 1152, 8])
        inputs = squash(reshaped)

        u_expanded_0 = tf.expand_dims(inputs, -1)
        u_expanded_1 = tf.expand_dims(u_expanded_0, 2)

        u_tiled = tf.tile(u_expanded_1, [1, 1, 10, 1, 1])
        W_tiled = tf.tile(self.kernel, [batch_size, 1, 1, 1, 1])

        u_hat = tf.matmul(W_tiled, u_tiled)
        try:
            numpy_result = np.matmul(W_tiled.numpy(), u_tiled.numpy())
            tf.print('\n TensorFlow/numpy max element value=' + str(tf.reduce_max(u_hat).numpy()) + '/' + str(
                numpy_result.max()))
            tf.print('\n TensorFlow/numpy mean value=' + str(tf.reduce_mean(u_hat).numpy()) + '/' + str(
                numpy_result.mean()))
        except:
            pass

        soft = softmax((safe_norm(tf.reduce_sum(u_hat, axis=[1, 3]))))

        # tf.print('\n\nW_tile max=' + str(tf.reduce_max(W_tiled)))
        # tf.print('W_tile min=' + str(tf.reduce_min(W_tiled)))
        # tf.print('u_tiled max=' + str(tf.reduce_max(u_tiled)))
        # tf.print('u_tiled min=' + str(tf.reduce_min(u_tiled)))
        # tf.print('u_hat max=' + str(tf.reduce_max(u_hat)))
        # tf.print('u_hat min=' + str(tf.reduce_min(u_hat)))

        tf.debugging.check_numerics(u_tiled, 'u_tiled')
        tf.debugging.check_numerics(W_tiled, 'W_tiled')
        tf.debugging.check_numerics(u_hat, 'u_hat')
        tf.debugging.check_numerics(soft, 'soft')

        return soft

    def compute_output_shape(self, input_shape):
        return (batch_size, 10, 16)


batch_size = 128
epochs = 100
img_rows, img_cols = 28, 28
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)

x_train = x_train.astype('float32')
x_train /= 255
y_train = utils.to_categorical(y_train, 10)

model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Conv2D(256, (9, 9), activation='relu', input_shape=(28, 28, 1)))
model.add(tf.keras.layers.Conv2D(256, (9, 9), strides=(2, 2), activation='relu'))
model.add(Custom_layer())

model.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])
model.run_eagerly = True
model.summary()

model.fit(
    x_train,
    y_train,
    batch_size=batch_size,
    epochs=epochs
)

我期望在 GPU 和 CPU 上得到类似的结果，但实际情况却大不相同。

【问题讨论】：

感谢您提供详细信息。你能确认每种情况下损失的dtype吗？我的预感是，不知何故，CPU 版本使用 64 位浮点运算，而 GPU 仅支持 32 位。

我已经编写了我的自定义损失函数：l = tf.keras.losses.categorical_crossentropy(y_true, y_pred); tf.print('Loss dtype:' + str(l.dtype)); return l 在标准输出中我看到：Loss dtype:<dtype: 'float32'>（CPU 和 GPU）。此外，我已经尝试（在 GPU 上）在 tf.matmul 之前将 u_tiled 和 W_tiled 转换为 tf.float64，并将产品转换为 tf.float32 - 它可以工作但速度很慢。

【参考方案1】：

在此处（回答部分）提供解决方案，即使它是由 user12292000 在问题中指定的，也是为了社区的利益。

官方 TF 2.0 build需要CUDA Toolkit 10.0（10.1不行），所以他重新编译了TF 2.0分支使用 CUDA 工具包 10.1。现在他在 TF 上的自定义版本按预期工作，is available here。

此外，在 official TF 2.0 构建（CUDA Toolkit 10.0）上，他复制了tf.matmul bug from official GitHub repo，但在他的自定义构建（CUDA Toolkit 10. 1) 也可以按预期工作。

【参考方案2】：

官方 TF 2.0 build需要CUDA Toolkit 10.0（10.1不行），所以我重新编译了TF 2.0分支使用 CUDA 工具包 10.1。我在 TF 上的自定义版本按预期工作，is available here。

此外，在 official TF 2.0 build (CUDA Toolkit 10.0) 我复制了 tf.matmul bug from official GitHub repo，但在我的自定义 build (CUDA Toolkit 10. 1) 也可以按预期工作。