【Python/Tensorflow】SequenceExampleを使用したSparseTensorのシリアライズ方法

#【TFReocrd作成側】
import numpy as np
import tensorflow as tf
import scipy

def int64_feature(l):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=l))

def float_feature(l):
    return tf.train.Feature(float_list=tf.train.FloatList(value=l))

# メイン
sample_mat = scipy.sparse.random(12, 12, density=0.1)

# 読み出すことを考えて、[行、列]の形式で保存
indices_features = []
for row, col in zip(sample_mat.row, sample_mat.col):
    indices_features.append(int64_feature([row, col]))

values_feature = float_feature(sample_mat.data)
shape_feature = int64_feature([12, 12])

feature_lists = tf.train.FeatureLists(feature_list={
    "indices": tf.train.FeatureList(feature=indices_features),
    "values":  tf.train.FeatureList(feature=[values_feature]),
    "shape":   tf.train.FeatureList(feature=[shape_feature])
})
seq_example = tf.train.SequenceExample(feature_lists=feature_lists)

savepath = "sparse_seq_example.tfrecords"
# TFRecord形式で保存
with tf.io.TFRecordWriter(savepath) as writer:
    writer.write(seq_example.SerializeToString())

print(seq_example)
"""
feature_lists {
  feature_list {
    key: "indices"
    value {
      feature {
        int64_list {
          value: 11
          value: 5
        }
      }
      - 長いので中略 -
      feature {
        int64_list {
          value: 9
          value: 11
        }
      }
    }
  }
  feature_list {
    key: "shape"
    value {
      feature {
        int64_list {
          value: 12
          value: 12
        }
      }
    }
  }
  feature_list {
    key: "values"
    value {
      feature {
        float_list {
          value: 0.4992492198944092
          value: 0.3927936255931854
          value: 0.23264065384864807
          value: 0.26840609312057495
          value: 0.7273346781730652
          value: 0.10040651261806488
          value: 0.8834704160690308
          value: 0.31499430537223816
          value: 0.25683218240737915
          value: 0.6979917287826538
          value: 0.7752628922462463
          value: 0.9083815217018127
          value: 0.7380165457725525
        }
      }
    }
  }
}
"""

# 【TFRecord読み出し側】
import numpy as np
import tensorflow as tf
import scipy

filepath = "test_sparse_seq_ex.tfrecords"

# shapeも含めて読み出せる点はこちらの方がSparseFeatureよりも有利か


def parse_seq_example(example):
    sequence_features = {
        "indices": tf.io.VarLenFeature(dtype=tf.int64),
        "values":  tf.io.VarLenFeature(dtype=tf.float32),
        "shape":   tf.io.VarLenFeature(dtype=tf.int64)
    }

    _, train_data = tf.io.parse_single_sequence_example(
        example,
        sequence_features=sequence_features
    )

    indices = tf.sparse.to_dense(train_data["indices"])
    values =  tf.reshape(tf.sparse.to_dense(train_data["values"]), shape=(-1,))
    shape =   tf.reshape(tf.sparse.to_dense(train_data["shape"]), shape=(-1,))

    train_X = tf.sparse.SparseTensor(indices, values, shape)

    return indices, values, shape


dataset = tf.data.TFRecordDataset([filepath]).map(parse_seq_example)

for data in dataset:
    print(data)

"""
【出力結果】
SparseTensor(
indices=tf.Tensor(
[[11  5]
 [ 9  2]
 [ 2 10]
 [ 9  6]
 [ 1  6]
 [11  1]
 [ 9 10]
 [10  4]
 [ 8 10]
 [ 3  0]
 [ 0  3]
 [ 5 11]
 [ 4  5]
 [ 9 11]], shape=(14, 2), dtype=int64),
values=tf.Tensor(
[0.49924922 0.39279363 0.23264065 0.2684061  0.7273347  0.10040651
 0.8834704  0.7617732  0.3149943  0.25683218 0.6979917  0.7752629
 0.9083815  0.73801655], shape=(14,), dtype=float32),
dense_shape=tf.Tensor([12 12], shape=(2,), dtype=int64))
"""