TF Recommenders, TF Agents, TF Federated

Three specialized stacks for three problem shapes

TF Recommenders (TFRS) provides building blocks for recommendation systems — the two-stage retrieval → ranking pipeline that powers Netflix, YouTube, Spotify. The retrieval stage uses ANN (approximate nearest neighbor) to narrow millions of candidates to thousands in milliseconds. The ranking stage applies an expensive model to score those candidates and pick the top-N.

TF Agents is a production-grade reinforcement learning library: DQN, PPO, SAC, TD3, DDPG. Discrete action spaces use DQN; continuous action spaces (robot torques, vehicle steering) use SAC.

TF Federated (TFF) enables ML on decentralized data. Training data stays on client devices (phones, hospitals); only model updates travel to the central server. The privacy-preserving alternative when raw data can't leave its origin.

Code

TFRS two-tower retrieval·python

import tensorflow as tf
import tensorflow_recommenders as tfrs

# pip install tensorflow-recommenders

user_model = tf.keras.Sequential([
    tf.keras.layers.StringLookup(vocabulary=unique_user_ids, mask_token=None),
    tf.keras.layers.Embedding(len(unique_user_ids) + 1, 64),
])
movie_model = tf.keras.Sequential([
    tf.keras.layers.StringLookup(vocabulary=unique_movie_titles, mask_token=None),
    tf.keras.layers.Embedding(len(unique_movie_titles) + 1, 64),
])

task = tfrs.tasks.Retrieval(
    metrics=tfrs.metrics.FactorizedTopK(
        candidates=movies.batch(128).map(movie_model),
    )
)

class MovielensModel(tfrs.Model):
    def compute_loss(self, features, training=False):
        u_emb = self.user_model(features["user_id"])
        m_emb = self.movie_model(features["movie_title"])
        return self.task(u_emb, m_emb)

model = MovielensModel(user_model, movie_model, task)
model.compile(optimizer=tf.keras.optimizers.Adagrad(0.5))
model.fit(ratings.batch(4096), epochs=3)

TF-Agents — DQN on CartPole·python

import tensorflow as tf
from tf_agents.environments import suite_gym, tf_py_environment
from tf_agents.networks import q_network
from tf_agents.agents.dqn import dqn_agent

# pip install tf-agents

env = suite_gym.load('CartPole-v1')
train_env = tf_py_environment.TFPyEnvironment(env)

q_net = q_network.QNetwork(
    train_env.observation_spec(),
    train_env.action_spec(),
    fc_layer_params=(128, 64),
)

agent = dqn_agent.DqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    q_network=q_net,
    optimizer=tf.keras.optimizers.Adam(1e-3),
    train_step_counter=tf.Variable(0),
)
agent.initialize()

TFF — federated averaging·python

import tensorflow_federated as tff

# pip install tensorflow-federated
trainer = tff.learning.algorithms.build_weighted_fed_avg(
    tff_model,
    client_optimizer_fn=tff.learning.optimizers.build_sgdm(learning_rate=0.1),
)

state = trainer.initialize()
for round_num in range(10):
    result = trainer.next(state, train_data)
    state  = result.state
    metrics = result.metrics
    print(f"Round {round_num}: "
          f"acc={metrics['client_work']['train']['accuracy']:.4f}")

TF Recommenders, TF Agents, TF Federated

Three specialized stacks for three problem shapes

Code

Progress

Comments 0