25 lines
798 B
Python
25 lines
798 B
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import wandb
|
|
from slime.utils.tensorboard_utils import _TensorboardAdapter
|
|
|
|
from . import wandb_utils
|
|
|
|
|
|
def init_tracking(args, primary: bool = True, **kwargs):
|
|
if primary:
|
|
wandb_utils.init_wandb_primary(args, **kwargs)
|
|
else:
|
|
wandb_utils.init_wandb_secondary(args, **kwargs)
|
|
|
|
|
|
# TODO further refactor, e.g. put TensorBoard init to the "init" part
|
|
def log(args, metrics, step_key: str):
|
|
if args.use_wandb:
|
|
wandb.log(metrics)
|
|
|
|
if args.use_tensorboard:
|
|
metrics_except_step = {k: v for k, v in metrics.items() if k != step_key}
|
|
_TensorboardAdapter(args).log(data=metrics_except_step, step=metrics[step_key])
|