如何将Scipy的csr_matrix转换为Tensorflow的SparseTensor并解决相关报错

问题原因 #

你的判断完全正确：scipy.sparse.csr_matrix的indices属性仅存储非零元素的列索引，是一维数组，而TensorFlow的SparseTensor要求indices参数为二维数组，形状为[N, 2]，其中N为非零元素总数，每行对应一个非零元素的[行索引, 列索引]，所以直接传入csr的indices会触发维度不匹配的报错。

解决方案 #

我们可以通过csr矩阵的indptr属性生成每个非零元素对应的行索引，再和列索引拼接得到符合要求的二维indices，代码如下：

import numpy as np
import tensorflow as tf
from scipy.sparse import load_npz
import os

# 读取npz中的csr矩阵
x_train_path = os.path.join('/home/gaurav.gupta/myproject', 'tmp', 'x_train.npz')
x_train = load_npz(x_train_path)

# 生成行索引：indptr[i]到indptr[i+1]之间的所有非零元素都属于第i行
row_indices = np.repeat(np.arange(x_train.shape[0]), np.diff(x_train.indptr))
# 拼接得到[N,2]的二维索引数组
indices = np.stack([row_indices, x_train.indices], axis=1)

# 构造SparseTensor
x_train_tf = tf.sparse.SparseTensor(
    indices=indices,
    values=x_train.data,
    dense_shape=x_train.shape
)

# 可选：如果后续操作要求索引按行主序排列，加上这行即可
x_train_tf = tf.sparse.reorder(x_train_tf)

补充说明 #

csr_matrix.indptr是长度为行数+1的数组，indptr[i]表示第i行第一个非零元素在data/indices数组中的起始位置，np.diff(x_train.indptr)得到每行的非零元素个数，用np.repeat就可以生成和非零元素一一对应的行索引。csr矩阵的非零元素本身就是按行优先存储的，拼接得到的indices默认符合行号递增的顺序，大部分场景下可以直接使用。

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