使用带有编码词汇表的 StringLookup 层保存张量流模型

Question

我在保存经过训练的 TensorFlow 模型时遇到了一些问题，我在该模型中有一个 StringLookup 层并且我需要使用 TFRecods 作为训练输入。重现问题的最小示例：

首先我定义训练数据

vocabulary = [str(i) for i in range(100, 200)]
X_train = np.random.choice(vocabulary, size=(100,))
y_train = np.random.choice([0,1], size=(100,))

我将它保存在文件中作为 tfrecords

def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def _string_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[str(value).encode('utf-8')]))

with tf.io.TFRecordWriter('train.tfrecords') as writer:
    for i in range(len(X_train)):
        example = tf.train.Example(features=tf.train.Features(feature={
            'user_id': _string_feature(X_train[i]),
            'label': _int64_feature(y_train[i])
        }))
        writer.write(example.SerializeToString())

然后我使用 tf.data API 能够将数据流式传输到训练中（原始数据不适合内存）

data = tf.data.TFRecordDataset(['train.tfrecords'])
features = {
    'user_id': tf.io.FixedLenFeature([], tf.string),
    'label': tf.io.FixedLenFeature([], tf.int64)
} 
def parse(record):
    parsed = tf.io.parse_single_example(record, features)
    return (parsed['user_id'], parsed['label'])
data = data.map(parse)

数据如下所示：

print(list(data.take(5).as_numpy_iterator()))
[(b'166', 1), (b'144', 0), (b'148', 1), (b'180', 0), (b'192', 0)]

原始数据集的字符串在此过程中被转换为字节。我必须将这个新词汇传递给 StringLookup 构造函数，因为传递字符串和使用字节进行训练会引发错误

new_vocab = [w.encode('utf-8') for w in vocabulary]

inp = tf.keras.Input(shape=(1,), dtype=tf.string)
x = tf.keras.layers.StringLookup(vocabulary=new_vocab)(inp)
x = tf.keras.layers.Embedding(len(new_vocab)+1, 32)(x)
out = tf.keras.layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.Model(inputs=[inp], outputs=[out])

model.compile(optimizer='adam', loss='BinaryCrossentropy')
model.fit(data.batch(10), epochs=5)

但是当我尝试保存模型时，出现错误，因为输入到 StringLookup 层的词汇表被编码为字节，无法转储到 json

model.save('model/')
TypeError: ('Not JSON Serializable:', b'100')

我真的不知道该怎么做，我读到 TensorFlow 建议使用编码字符串而不是普通字符串，但这不允许保存模型。我还尝试在将字符串馈送到模型之前对数据进行解码，但如果不将所有数据加载到内存中（仅使用 tf.data 操作）我无法做到这一点

Answer 1

使用您的数据和原始词汇：

import tensorflow as tf
import numpy as np

vocabulary = [str(i) for i in range(100, 200)]
X_train = np.random.choice(vocabulary, size=(100,))
y_train = np.random.choice([0,1], size=(100,))
...
...
data = data.map(parse)

我运行你的代码（有一个额外的 Flatten 层）并且能够保存你的模型：

inp = tf.keras.Input(shape=(1,), dtype=tf.string)
x = tf.keras.layers.StringLookup(vocabulary=vocabulary)(inp)
x = tf.keras.layers.Embedding(len(vocabulary)+1, 32)(x)
x = tf.keras.layers.Flatten()(x)
out = tf.keras.layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.Model(inputs=[inp], outputs=[out])

model.compile(optimizer='adam', loss='BinaryCrossentropy')
model.fit(data.batch(10), epochs=5)
model.save('model/')

Epoch 1/5
10/10 [==============================] - 1s 8ms/step - loss: 0.6949
Epoch 2/5
10/10 [==============================] - 0s 4ms/step - loss: 0.6864
Epoch 3/5
10/10 [==============================] - 0s 5ms/step - loss: 0.6787
Epoch 4/5
10/10 [==============================] - 0s 5ms/step - loss: 0.6707
Epoch 5/5
10/10 [==============================] - 0s 5ms/step - loss: 0.6620
INFO:tensorflow:Assets written to: model/assets

我不明白你为什么需要 new_vocab = [w.encode('utf-8') for w in vocabulary]。

如果您确实需要使用 new_vocab，您可以尝试在训练期间设置它，然后设置 vocabulary 来保存您的模型，因为唯一的区别是编码：

new_vocab = [w.encode('utf-8') for w in vocabulary]

lookup_layer = tf.keras.layers.StringLookup()
lookup_layer.adapt(new_vocab)
inp = tf.keras.Input(shape=(1,), dtype=tf.string)
x = lookup_layer(inp)
x = tf.keras.layers.Embedding(len(new_vocab)+1, 32)(x)
x = tf.keras.layers.Flatten()(x)
out = tf.keras.layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.Model(inputs=[inp], outputs=[out])

model.compile(optimizer='adam', loss='BinaryCrossentropy')
model.fit(data.batch(10), epochs=5)
model.layers[1].adapt(vocabulary)

model.save('/model')

诚然，这很老套。

使用带有编码词汇表的 StringLookup 层保存张量流模型

Save tensorflow model with StringLookup layer with encoded vocabulary

python

deep-learning

keras

tensorflow

tfrecord