挑选出tensor中等于0的索引_【tf.matmul 致命错误】请谨慎使用tensorflow 2.0

2020/1/11更新：

在 tensorflow 2.1 及以上版本中，该bug已解决。

如何升级到 2.1 及以上版本，请移步：

Kevin：Anaconda 搭建 Tensorflow 2 开发环境​zhuanlan.zhihu.com

---

A fatal error is a bug that does not report any errors, so we cannot locate or resolve it.

已经有许多博文介绍了tensorflow 2.0先进的设计理念和人性化的新特性：相较于tf1.0，tf2.0 删去了很多反直觉的概念和方法，并且使用动态图机制，与python完美融合等等等。而且近半年tf2.0版本不断迭代，终于从测试版的α、β、c一路更新到稳定版的stable，是时候尽情拥抱tf2.0了......了吗？

稳定版 tf2.0 stable 真的 stable 吗？

tensorflow 是一个用于数值计算的库（废话），矩阵运算是它进行数值计算的基础操作（废话），那么如果矩阵运算，即常用的tf.matmul操作出现了一个诡异的、不会报错的bug呢？ (;° ロ°)

github 上关于这个 bug 的 issue：A puzzling & fatal error occurred in the tf.matmal()

该 bug 在colab上的复现：here

高维张量的乘法

首先再复习一下，高维张量的乘法：

'''
对于两个张量：
a.shape = [dim_1,...,dim_n, l, k]
b.shape = [dim_1,...,dim_n, k, m]
只要最后的两个维度满足矩阵乘法的要求，而其他维度（并行维度）大小相等，则可以进行如下乘法操作：
'''
c = tf.matmul(a, b)
'''
结果：
c.shape = [dim_1,...,dim_n, l, m]
'''

展开来相当于并行计算多组矩阵乘法，因此这种操作非常适合于在gpu上运行：

显然地，并行计算的结果应该等于分开来计算的结果，亦即tf.matmul(a[i], b[i])应该等于tf.matmul(a, b)[i]。但是当tensor的维度较大时，这个显然的性质竟然不成立，bug来了。

bug初体验

你可以通过下面的例子来体会这个bug：

当使用gpu计算较大维度tensor的tf.matmul时，将会出现如下错误：

j = np.random.rand(10, 6, 1130, 16, 8)
k = np.random.rand(10, 6, 1130, 8, 1)

j = tf.cast(j, dtype=tf.float32)
k = tf.cast(k, dtype=tf.float32)

a = tf.matmul(j, k)[9, 3]  # 大维度tensor相乘
b = tf.matmul(j[9], k[9])[3]
c = tf.matmul(j[9, 3], k[9, 3])

print(tf.reduce_all(tf.equal(a, b)))
print(tf.reduce_sum(a-b))
print(tf.reduce_all(tf.equal(b, c)))

'''
tf.Tensor(False, shape=(), dtype=bool)  # 正确的值应该是 True
tf.Tensor(-1.3804454e+38, shape=(), dtype=float32)  # 比较了a和b的具体差异，发现造成错误的并不是细微的差异
tf.Tensor(True, shape=(), dtype=bool)
'''

而使用cpu则不会出现该错误：

...

with tf.device("CPU:0"):

    a = tf.matmul(j, k)[9, 3]
    b = tf.matmul(j[9], k[9])[3]
    c = tf.matmul(j[9, 3], k[9, 3])

    print(tf.reduce_all(tf.equal(a, b)))
    print(tf.reduce_sum(a-b))
    print(tf.reduce_all(tf.equal(b, c)))

'''
tf.Tensor(True, shape=(), dtype=bool)
tf.Tensor(0.0, shape=(), dtype=float32)
tf.Tensor(True, shape=(), dtype=bool)
'''

有趣的是，只要稍微将其中一个维度变小，比如将 1130 减一，同样使用gpu，也不会发生错误。

# j = np.random.rand(10, 6, 1130, 16, 8)
# k = np.random.rand(10, 6, 1130, 8, 1)
j = np.random.rand(10, 6, 1129, 16, 8)  # 1130 --> 1129
k = np.random.rand(10, 6, 1129, 8, 1)

j = tf.cast(j, dtype=tf.float32)
k = tf.cast(k, dtype=tf.float32)

a = tf.matmul(j, k)[9, 3]
b = tf.matmul(j[9], k[9])[3]
c = tf.matmul(j[9, 3], k[9, 3])

print(tf.reduce_all(tf.equal(a, b)))
print(tf.reduce_sum(a-b))
print(tf.reduce_all(tf.equal(b, c)))

'''
tf.Tensor(True, shape=(), dtype=bool)
tf.Tensor(0.0, shape=(), dtype=float32)
tf.Tensor(True, shape=(), dtype=bool)
'''

进一步地，我们猜想对于使用gpu的tf.matmul操作，输入tensor的并行维度有一个上限大小，超过它就会出现bug。我进行了如下的测试，发现对于

的矩阵相乘，并行维度超过

将会产生bug。

offset = 0
while True:
    j = np.random.rand(*(65530+offset*1 , 16, 8))
    k = np.random.rand(*(65530+offset*1 , 8, 1))
    # with tf.device("CPU:0"):
    j = tf.cast(j, dtype=tf.float32)
    k = tf.cast(k, dtype=tf.float32)

    a = tf.matmul(j, k)[-1]

    b = tf.matmul(j[-1], k[-1])
    print(offset)

    if not tf.reduce_all(tf.equal(a, b)).numpy():
        break

    offset += 1

print(65530+offset*1)
'''
65536  (is 2^16)
'''

但是很遗憾的是，对于不同大小的矩阵相乘，并行维度的上限并不确定，似乎是矩阵越大、并行维度上限越小，但并没有确切的规律。

bug特点

首先：虽然我们知道并行维度过大将会产生bug，但是我们并不知道过大的标准是什么，而且这个标准似乎随着矩阵的大小在不规律变化。

并且：这个 bug 不会报错，你不知道它是否发生、什么时候发生。

因此：这是一个不可控的bug（除非你不使用gpu，但那样你为什么还要用tensorflow呢）。

tf.matmul的应用广泛，比如卷积的底层实现就用到了，因此这种bug会影响网络的深层运行，你又一时半会发现不了，排查也非常困难。

其他

这个bug是否产生自显卡内存的限制？

不是，无论是在仅有2gb内存的gtx 850，还是8g内存的rtx 2070，bug都在矩阵大小超过相同值时产生。

bug与系统有关吗？

无论是 Linux Ubuntu 18.04 还是 win 10 都存在这个bug。

bug与tf2.0的版本有关吗？

在tf2.0 α、β、c 和 stable 上都存在这个bug。