内存高效的推荐系统嵌入

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

from zipfile import ZipFile
from urllib.request import urlretrieve
import numpy as np
import pandas as pd
import tensorflow as tf
import keras
from keras import layers
from keras.layers import StringLookup
import matplotlib.pyplot as plt

urlretrieve("http://files.grouplens.org/datasets/movielens/ml-1m.zip", "movielens.zip")
ZipFile("movielens.zip", "r").extractall()

ratings_data = pd.read_csv(
    "ml-1m/ratings.dat",
    sep="::",
    names=["user_id", "movie_id", "rating", "unix_timestamp"],
)

ratings_data["movie_id"] = ratings_data["movie_id"].apply(lambda x: f"movie_{x}")
ratings_data["user_id"] = ratings_data["user_id"].apply(lambda x: f"user_{x}")
ratings_data["rating"] = ratings_data["rating"].apply(lambda x: float(x))
del ratings_data["unix_timestamp"]

print(f"用户数量: {len(ratings_data.user_id.unique())}")
print(f"电影数量: {len(ratings_data.movie_id.unique())}")
print(f"评分数量: {len(ratings_data.index)}")

random_selection = np.random.rand(len(ratings_data.index)) <= 0.85
train_data = ratings_data[random_selection]
eval_data = ratings_data[~random_selection]

train_data.to_csv("train_data.csv", index=False, sep="|", header=False)
eval_data.to_csv("eval_data.csv", index=False, sep="|", header=False)
print(f"训练数据分割: {len(train_data.index)}")
print(f"验证数据分割: {len(eval_data.index)}")
print("训练和验证数据文件已保存。")

csv_header = list(ratings_data.columns)
user_vocabulary = list(ratings_data.user_id.unique())
movie_vocabulary = list(ratings_data.movie_id.unique())
target_feature_name = "rating"
learning_rate = 0.001
batch_size = 128
num_epochs = 3
base_embedding_dim = 64

def get_dataset_from_csv(csv_file_path, batch_size=128, shuffle=True):
    return tf.data.experimental.make_csv_dataset(
        csv_file_path,
        batch_size=batch_size,
        column_names=csv_header,
        label_name=target_feature_name,
        num_epochs=1,
        header=False,
        field_delim="|",
        shuffle=shuffle,
    )


def run_experiment(model):
    # 编译模型。
    model.compile(
        optimizer=keras.optimizers.Adam(learning_rate),
        loss=keras.losses.MeanSquaredError(),
        metrics=[keras.metrics.MeanAbsoluteError(name="mae")],
    )
    # 读取训练数据。
    train_dataset = get_dataset_from_csv("train_data.csv", batch_size)
    # 读取测试数据。
    eval_dataset = get_dataset_from_csv("eval_data.csv", batch_size, shuffle=False)
    # 用训练数据拟合模型。
    history = model.fit(
        train_dataset,
        epochs=num_epochs,
        validation_data=eval_dataset,
    )
    return history

def embedding_encoder(vocabulary, embedding_dim, num_oov_indices=0, name=None):
    return keras.Sequential(
        [
            StringLookup(
                vocabulary=vocabulary, mask_token=None, num_oov_indices=num_oov_indices
            ),
            layers.Embedding(
                input_dim=len(vocabulary) + num_oov_indices, output_dim=embedding_dim
            ),
        ],
        name=f"{name}_embedding" if name else None,
    )

def create_baseline_model():
    # 接收用户作为输入。
    user_input = layers.Input(name="user_id", shape=(), dtype=tf.string)
    # 获取用户嵌入。
    user_embedding = embedding_encoder(
        vocabulary=user_vocabulary, embedding_dim=base_embedding_dim, name="user"
    )(user_input)

    # 接收电影作为输入。
    movie_input = layers.Input(name="movie_id", shape=(), dtype=tf.string)
    # 获取嵌入。
    movie_embedding = embedding_encoder(
        vocabulary=movie_vocabulary, embedding_dim=base_embedding_dim, name="movie"
    )(movie_input)

    # 计算用户与电影嵌入之间的点积相似性。
    logits = layers.Dot(axes=1, name="dot_similarity")(
        [user_embedding, movie_embedding]
    )
    # 转换为评分尺度。
    prediction = keras.activations.sigmoid(logits) * 5
    # 创建模型。
    model = keras.Model(
        inputs=[user_input, movie_input], outputs=prediction, name="baseline_model"
    )
    return model


baseline_model = create_baseline_model()
baseline_model.summary()

history = run_experiment(baseline_model)

plt.plot(history.history["loss"])
plt.plot(history.history["val_loss"])
plt.title("模型损失")
plt.ylabel("损失")
plt.xlabel("时期")
plt.legend(["训练", "评估"], loc="upper left")
plt.show()

class QREmbedding(keras.layers.Layer):
    def __init__(self, vocabulary, embedding_dim, num_buckets, name=None):
        super().__init__(name=name)
        self.num_buckets = num_buckets

        self.index_lookup = StringLookup(
            vocabulary=vocabulary, mask_token=None, num_oov_indices=0
        )
        self.q_embeddings = layers.Embedding(
            num_buckets,
            embedding_dim,
        )
        self.r_embeddings = layers.Embedding(
            num_buckets,
            embedding_dim,
        )

    def call(self, inputs):
        # 获取项目索引。
        embedding_index = self.index_lookup(inputs)
        # 获取商索引。
        quotient_index = tf.math.floordiv(embedding_index, self.num_buckets)
        # 获取余数索引。
        remainder_index = tf.math.floormod(embedding_index, self.num_buckets)
        # 使用商索引查找 quotient_embedding。
        quotient_embedding = self.q_embeddings(quotient_index)
        # 使用余数索引查找 remainder_embedding。
        remainder_embedding = self.r_embeddings(remainder_index)
        # 使用乘法作为组合操作
        return quotient_embedding * remainder_embedding

class MDEmbedding(keras.layers.Layer):
    def __init__(
        self, blocks_vocabulary, blocks_embedding_dims, base_embedding_dim, name=None
    ):
        super().__init__(name=name)
        self.num_blocks = len(blocks_vocabulary)

        # 创建词汇到块的查找。
        keys = []
        values = []
        for block_idx, block_vocab in enumerate(blocks_vocabulary):
            keys.extend(block_vocab)
            values.extend([block_idx] * len(block_vocab))
        self.vocab_to_block = tf.lookup.StaticHashTable(
            tf.lookup.KeyValueTensorInitializer(keys, values), default_value=-1
        )

        self.block_embedding_encoders = []
        self.block_embedding_projectors = []

        # 创建块嵌入编码器和投影器。
        for idx in range(self.num_blocks):
            vocabulary = blocks_vocabulary[idx]
            embedding_dim = blocks_embedding_dims[idx]
            block_embedding_encoder = embedding_encoder(
                vocabulary, embedding_dim, num_oov_indices=1
            )
            self.block_embedding_encoders.append(block_embedding_encoder)
            if embedding_dim == base_embedding_dim:
                self.block_embedding_projectors.append(layers.Lambda(lambda x: x))
            else:
                self.block_embedding_projectors.append(
                    layers.Dense(units=base_embedding_dim)
                )

    def call(self, inputs):
        # 获取每个输入项的块索引。
        block_indicies = self.vocab_to_block.lookup(inputs)
        # 将输出嵌入初始化为零。
        embeddings = tf.zeros(shape=(tf.shape(inputs)[0], base_embedding_dim))
        # 从块生成嵌入。
        for idx in range(self.num_blocks):
            # 查找当前块的嵌入。
            block_embeddings = self.block_embedding_encoders[idx](inputs)
            # 将嵌入投影到 base_embedding_dim。
            block_embeddings = self.block_embedding_projectors[idx](block_embeddings)
            # 创建一个掩码以过滤掉不属于当前块的项的嵌入。
            mask = tf.expand_dims(tf.cast(block_indicies == idx, tf.dtypes.float32), 1)
            # 将不属于当前块的项的嵌入设置为零。
            block_embeddings = block_embeddings * mask
            # 将块嵌入添加到最终嵌入中。
            embeddings += block_embeddings

        return embeddings

movie_frequencies = ratings_data["movie_id"].value_counts()
movie_frequencies.hist(bins=10)

sorted_movie_vocabulary = list(movie_frequencies.keys())

movie_blocks_vocabulary = [
    sorted_movie_vocabulary[:400],  # 高人气电影块
    sorted_movie_vocabulary[400:1700],  # 普通人气电影块
    sorted_movie_vocabulary[1700:],  # 低人气电影块
]

movie_blocks_embedding_dims = [64, 32, 16]

user_embedding_num_buckets = len(user_vocabulary) // 50


def create_memory_efficient_model():
    # 将用户作为输入。
    user_input = layers.Input(name="user_id", shape=(), dtype="string")
    # 获取用户嵌入。
    user_embedding = QREmbedding(
        vocabulary=user_vocabulary,
        embedding_dim=base_embedding_dim,
        num_buckets=user_embedding_num_buckets,
        name="user_embedding",
    )(user_input)

    # 将电影作为输入。
    movie_input = layers.Input(name="movie_id", shape=(), dtype="string")
    # 获取嵌入。
    movie_embedding = MDEmbedding(
        blocks_vocabulary=movie_blocks_vocabulary,
        blocks_embedding_dims=movie_blocks_embedding_dims,
        base_embedding_dim=base_embedding_dim,
        name="movie_embedding",
    )(movie_input)

    # 计算用户和电影嵌入之间的点积相似度。
    logits = layers.Dot(axes=1, name="dot_similarity")(
        [user_embedding, movie_embedding]
    )
    # 转换为评分尺度。
    prediction = keras.activations.sigmoid(logits) * 5
    # 创建模型。
    model = keras.Model(
        inputs=[user_input, movie_input], outputs=prediction, name="baseline_model"
    )
    return model


memory_efficient_model = create_memory_efficient_model()
memory_efficient_model.summary()

history = run_experiment(memory_efficient_model)

plt.plot(history.history["loss"])
plt.plot(history.history["val_loss"])
plt.title("模型损失")
plt.ylabel("损失")
plt.xlabel("纪元")
plt.legend(["训练", "评估"], loc="upper left")
plt.show()

Epoch 1/3
   6622/未知  6s 891微秒/步骤 - 损失: 1.1938 - 平均绝对误差: 0.8780

/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/contextlib.py:153: 用户警告: 您的输入数据已耗尽；中断训练。确保您的数据集或生成器可以生成至少 `steps_per_epoch * epochs` 批次。您可能需要在构建数据集时使用 `.repeat()` 函数。
  self.gen.throw(typ, value, traceback)

 6646/6646 ━━━━━━━━━━━━━━━━━━━━ 7s 992微秒/步骤 - 损失: 1.1931 - 平均绝对误差: 0.8777 - 验证损失: 1.1027 - 验证平均绝对误差: 0.8179
Epoch 2/3
 6646/6646 ━━━━━━━━━━━━━━━━━━━━ 7s 1毫秒/步骤 - 损失: 0.8908 - 平均绝对误差: 0.7488 - 验证损失: 0.9144 - 验证平均绝对误差: 0.7549
Epoch 3/3
 6646/6646 ━━━━━━━━━━━━━━━━━━━━ 7s 980微秒/步骤 - 损失: 0.8419 - 平均绝对误差: 0.7278 - 验证损失: 0.8806 - 验证平均绝对误差: 0.7419