【Python/Tensorflow】SparseTensorを学習に用いる際の注意点

import tensorflow as tf
import numpy as np

# Sparse Tensor の情報
rows = [1, 2]
cols = [2, 3]
values = [1.1, 2.2]
shape = (4, 4)

# Sparse Tensor
sparse_tensor = tf.sparse.SparseTensor(indices=[rows, cols], values=values, dense_shape=shape)
print(sparse_tensor)
"""
【出力結果】
SparseTensor(
indices=tf.Tensor([[1 2]
                   [2 3]], shape=(2, 2), dtype=int64),
values=tf.Tensor([1.1 2.2], shape=(2,), dtype=float32),
dense_shape=tf.Tensor([4 4], shape=(2,), dtype=int64)
)
"""
# Dense Tensor
dense_tensor = tf.sparse.to_dense(sparse_tensor)
print(dense_tensor)
"""
【出力結果】
tf.Tensor(
[[0.  0.  0.  0. ]
 [0.  0.  1.1 0. ]
 [0.  0.  0.  2.2]
 [0.  0.  0.  0. ]], shape=(4, 4), dtype=float32)
"""

import numpy as np
import tensorflow as tf
import scipy

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

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

# Protocol Buffersの形式に変換
# ⇒ あとで3階テンソルとして扱う必要があるので、"channels"の要素も用意
def sp_tensor2example(sp_array, label):
    return tf.train.Example(features=tf.train.Features(feature={
        "rows": feature_int64_list(sp_array.row),
        "cols": feature_int64_list(sp_array.col),
        "channels": feature_int64_list(np.zeros(len(sp_array.row)).astype(int)),
        "values": feature_float_list(sp_array.data),
        "label": feature_float_list([label])
    }))

filepath = "./sequence_conv.tfrecords"

# 疎行列・正解データを10個ずつ作成
sp_array_list = [scipy.sparse.random(12, 12, density=0.1) for i in range(10)]
label_list = np.random.rand(10)

# TFRecords形式で保存
with tf.io.TFRecordWriter(filepath) as writer:
    for sp_array, label in zip(sp_array_list, label_list):
        example_proto = sp_tensor2example(sp_array, label)
        writer.write(example_proto.SerializeToString())

import numpy as np
import tensorflow as tf

# 保存しているSparseTensorの形状 ※ Conv2Dでは3階テンソルを要求されるのでこの形状
sparse_shape = [12, 12, 1]

# デシリアライズ用コールバック関数
def parse_example(example):
    features = {"train": tf.io.SparseFeature(index_key=["rows", "cols", "channels"],
                                             value_key="values",
                                             dtype=tf.float32,
                                             size=sparse_shape),
                "label": tf.io.FixedLenFeature([], dtype=tf.float32)
                }
    train_data = tf.io.parse_single_example(example, features=features)
    x = train_data["train"]
    y = train_data["label"]
    return x, y

# メイン
filepath = "./sequence_conv.tfrecords"
# TFRecords読み出し & デシリアライズ
dataset_train = tf.data.TFRecordDataset([filepath]).map(parse_example)

for train_X, train_y in dataset_train:
    print(train_X)
    print(train_y)

"""
【出力結果(長いので一部だけ)】
SparseTensor(
indices=tf.Tensor(
        [[ 0  0  0]
         [ 0  3  0]
         [ 0  6  0]
         [ 1  7  0]
         [ 2  9  0]
         [ 3  3  0]
         [ 3  8  0]
         [ 4 11  0]
         [ 6  6  0]
         [ 6  8  0]
         [ 8  8  0]
         [ 8 10  0]
         [11  2  0]
         [11 10  0]], shape=(14, 3), dtype=int64),
values=tf.Tensor(
       [0.7976593  0.61596143 0.8332581  0.83961755 0.69140923 0.5217606
        0.18110311 0.98035353 0.11858348 0.97409976 0.67042786 0.7941748
        0.27434942 0.6197932 ], shape=(14,), dtype=float32),
dense_shape=tf.Tensor([12 12  1], shape=(3,), dtype=int64)
)

tf.Tensor(0.4406326, shape=(), dtype=float32)
"""

import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense, Flatten, Lambda, Conv2D
from tensorflow.keras.optimizers import RMSprop

# 保存しているSparseTensorの形状 ※ Conv2Dでは3階テンソルを要求されるのでこの形状
sparse_shape = [12, 12, 1]

# 各種パラメータ
DATA_SIZE = 10
batch_size = 2
steps_per_epoch = (DATA_SIZE-1) // batch_size + 1
kernel_size = (4, 4)
num_filters = 4

# デシリアライズ用コールバック関数
def parse_example(example):
    features = {"train": tf.io.SparseFeature(index_key=["rows", "cols", "channels"],
                                             value_key="values",
                                             dtype=tf.float32,
                                             size=sparse_shape),
                "label": tf.io.FixedLenFeature([], dtype=tf.float32)
                }
    train_data = tf.io.parse_single_example(example, features=features)
    x = train_data["train"]
    y = train_data["label"]
    return x, y

# メイン
if __name__ == "__main__":
    # TFRecordsファイルへのパス
    filepath = "./sequence_conv.tfrecords"

    # TFRecordsファイルの読み出し ※学習用に、バッチサイズと繰り返し回数を指定
    dataset_train = tf.data.TFRecordDataset([filepath]) \
        .map(parse_example) \
        .batch(batch_size) \
        .repeat(-1)

    # Keras API でのモデル構築
    inputs = Input(batch_size=batch_size, shape=sparse_shape, sparse=True)
    X = Lambda(tf.sparse.to_dense)(inputs)
    X = Conv2D(filters=num_filters, kernel_size=kernel_size)(X)
    X = Flatten()(X)
    X = Dense(1, activation="relu")(X)

    model = Model(inputs, X)
    model.compile(
        loss="MSE",
        optimizer=RMSprop()
    )
    print(model.summary())

    # 学習実行
    model.fit(
        x=dataset_train,
        epochs=3,
        verbose=1,
        steps_per_epoch=steps_per_epoch
    )
"""
【出力結果】
Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
input_1 (InputLayer)         [(2, 12, 12, 1)]          0
_________________________________________________________________
lambda (Lambda)              (2, 12, 12, 1)            0
_________________________________________________________________
conv2d (Conv2D)              (2, 9, 9, 4)              68
_________________________________________________________________
flatten (Flatten)            (2, 324)                  0
_________________________________________________________________
dense (Dense)                (2, 1)                    325
=================================================================
Total params: 393
Trainable params: 393
Non-trainable params: 0
_________________________________________________________________
None
Train for 5 steps
Epoch 1/3
5/5 [==============================] - 0s 90ms/step - loss: 0.1210
Epoch 2/3
5/5 [==============================] - 0s 3ms/step - loss: 0.1038
Epoch 3/3
5/5 [==============================] - 0s 2ms/step - loss: 0.0836
"""