如何手动向Ray Train内置TensorBoard Logger写入日志？

1. 在训练循环内获取TBXLogger实例 #

Ray Train会在每个训练worker的进程中自动初始化TBXLogger，你只需在训练函数（比如传给TorchTrainer/TensorFlowTrainer的train_loop_per_worker）内部，通过ray.train.get_context().get_logger()获取该实例，进而调用TensorBoard的原生API记录自定义内容。

示例代码：

import ray.train as train
from ray.train.torch import TorchTrainer
import torch

def train_loop_per_worker():
    # 获取Ray Train内置的TBXLogger
    logger = train.get_context().get_logger()
    
    # 确认Logger类型并获取TensorBoard的SummaryWriter
    if isinstance(logger, train.TBXLogger):
        tb_writer = logger.writer
        
        # 记录自定义标量
        tb_writer.add_scalar("custom/train_accuracy", 0.85, global_step=1)
        
        # 记录自定义直方图（示例：模型权重分布）
        dummy_weights = torch.randn(256)
        tb_writer.add_histogram("custom/layer_weights", dummy_weights, global_step=1)
        
        # 记录自定义图像（CHW格式示例）
        dummy_image = torch.randn(3, 128, 128)
        tb_writer.add_image("custom/sample_input", dummy_image, global_step=1)
        
        # 记录PR曲线
        predictions = torch.randn(200)
        ground_truth = torch.randint(0, 2, (200,))
        tb_writer.add_pr_curve("custom/class_pr_curve", ground_truth, predictions, global_step=1)
    
    # 常规训练流程
    for epoch in range(10):
        # 模拟训练步骤
        train.report({"epoch_loss": 0.2 * (10 - epoch)})

# 启动训练
trainer = TorchTrainer(
    train_loop_per_worker,
    scaling_config=train.ScalingConfig(num_workers=2)
)
trainer.fit()

2. 关键注意事项 #

• 调用时机：必须在训练worker的循环函数内部调用get_logger()，外部初始化阶段调用会返回None。
• 多worker兼容：每个worker的TBXLogger会自动将日志写入统一的TensorBoard目录，无需手动合并。
• API兼容性：通过logger.writer获取的SummaryWriter和原生TensorBoard的API完全一致，所有原生支持的日志类型都可以直接使用。

内容的提问来源于stack exchange，提问作者Ricardo Decal