如何通过REST API监控TensorFlow、PyTorch等神经网络模型训练状态？

一、直接通过训练框架的回调/钩子对接API #

这是最直接的方式，利用框架原生的回调机制，在训练的关键节点（epoch结束、batch结束）将指标推送到REST API后端。

TensorFlow/Keras 实现示例 #

from tensorflow.keras.callbacks import Callback
import requests

class APIMonitorCallback(Callback):
    def __init__(self, api_endpoint, model_unique_id):
        self.api_endpoint = api_endpoint
        self.model_id = model_unique_id

    def on_epoch_end(self, epoch, logs=None):
        # 整理当前epoch的训练指标
        metric_data = {
            "model_id": self.model_id,
            "epoch": epoch,
            "train_loss": logs.get("loss"),
            "val_loss": logs.get("val_loss"),
            "train_acc": logs.get("accuracy"),
            "val_acc": logs.get("val_accuracy")
        }
        # 异步或同步推送至API后端（建议加超时避免阻塞训练）
        requests.post(f"{self.api_endpoint}/submit-metrics", json=metric_data, timeout=5)

# 训练时挂载回调
model.fit(train_data, epochs=50, callbacks=[APIMonitorCallback("http://your-api:5000", "resnet_101_cifar10")])

PyTorch 实现示例 #

PyTorch没有原生的全局回调，可在自定义训练循环中嵌入指标推送逻辑：

import requests

# 自定义训练循环
for epoch in range(total_epochs):
    train_loss, train_acc = run_train_epoch(model, train_loader)
    val_loss, val_acc = run_val_epoch(model, val_loader)

    # 推送指标到API
    metric_data = {
        "model_id": "vit_imagenet",
        "epoch": epoch,
        "train_loss": train_loss,
        "val_loss": val_loss,
        "train_acc": train_acc,
        "val_acc": val_acc
    }
    requests.post("http://your-api:5000/submit-metrics", json=metric_data, timeout=5)

二、用消息队列实现解耦（更优雅的方案） #

如果担心API服务不可用阻塞训练进程，或者需要同时对接多个监控服务，用消息队列做中间层实现解耦：

• 训练进程将指标推送到消息队列（如Redis Pub/Sub、RabbitMQ）
• REST API后端作为消费者，从队列拉取指标并存入内存或数据库
• 前端通过WebSocket实时订阅或轮询API获取指标

Redis Pub/Sub 示例（PyTorch + Flask） #

训练端推送指标

import redis
import json

redis_client = redis.Redis(host="localhost", port=6379, db=0)

# 训练循环内推送
metric_data = {"model_id": "lstm_text_class", "epoch": epoch, ...}
redis_client.publish("training_metrics", json.dumps(metric_data))

API后端消费并提供接口

from flask import Flask, jsonify
import redis
import json
from threading import Thread

app = Flask(__name__)
# 内存临时存储最新指标，历史数据可存入数据库
metrics_cache = {}

def subscribe_to_metrics():
    redis_client = redis.Redis(host="localhost", port=6379, db=0)
    pubsub = redis_client.pubsub()
    pubsub.subscribe("training_metrics")
    for msg in pubsub.listen():
        if msg["type"] == "message":
            data = json.loads(msg["data"])
            model_id = data["model_id"]
            if model_id not in metrics_cache:
                metrics_cache[model_id] = []
            metrics_cache[model_id].append(data)

# 启动后台订阅线程
Thread(target=subscribe_to_metrics, daemon=True).start()

# 给前端提供的指标查询接口
@app.route("/metrics/<model_id>")
def get_model_metrics(model_id):
    return jsonify(metrics_cache.get(model_id, []))

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)

三、复用现有监控工具的API #

既然你已经在使用TensorBoard或Weights & Biases，可以直接复用它们的API来获取指标：

• TensorBoard：可解析其日志目录下的事件文件，或调用其内置REST接口（默认http://localhost:6006/data/plugin/scalars/scalars）拉取标量指标
• Weights & Biases：用官方Python API在后端拉取训练数据，再转发给前端

四、关键注意事项 #

• 避免训练阻塞：直接调用API时要设置超时，或用异步请求库（如aiohttp）
• 持久化：内存存储重启会丢失，建议同时将指标存入数据库（如PostgreSQL、SQLite）
• 唯一标识：给每个训练任务分配唯一model_id，方便前端区分不同模型的指标

内容的提问来源于stack exchange，提问作者Dmitry Sokolov