我正在尝试减小两个稀疏张量的尺寸,将它们相乘(一个需要广播),然后再次减小reduce_sum。甚至在减少之后,至少有一个稀疏矩阵确实是稀疏的,并且我在GPU内存错误上失败。
tf.sparse_reduce_sum不能产生具有梯度的稀疏张量吗?可以有效地创建它吗?
最小代码:
import tensorflow as tf
from tensorflow.python.ops import math_ops
# outputs sparse dim = batch_dim, dim1, dim2, dim3, dim4
outputs = tf.sparse.reduce_sum(outputs, axis=-1)
# gates sparse dim = batch, dim1, dim2
gates = tf.sparse.reduce_sum(gates, axis=-1)
gates = math_ops.sigmoid(out)
gates = expand_dims(gates, -1)
gates = tf.tile(gates, [1, 1, 1, tf.shape(outputs)[-1]
outputs = math_ops.multiply(gates, outputs)
outputs = math_ops.reduce_sum(outputs, [-2])