from pathlib import Path
from typing import Optional, Union
from torch.utils.tensorboard import SummaryWriter
import wandb
[docs]
def get_writer(
logdir: Union[Path, str],
wandb_project: Optional[str] = None,
wandb_name: Optional[str] = None,
id: Optional[str] = None,
):
"""Get a TensorBoard writer for logging.
If a project name is given, also log to Weights & Biases.
Args:
logdir: The path to the directory where to store the
Weights & Biases and TensorBoard logs. If None, no logs are written.
wandb_project: The name of the Weights & Biases project to log to.
If None, no logs are written to Weights & Biases.
wandb_name: The name of the run in Weights & Biases. If None, a two-word name
is generated by wandb.
id: The id of the run in Weights & Biases (that will be continued). If None, a
new run is created.
Returns:
A TensorBoard writer or None if logdir is None.
"""
if logdir is None:
writer = None
else:
if wandb_project is not None:
# make directory for wandb logs
if id is not None:
wandb.init(
project=wandb_project,
sync_tensorboard=True,
dir=logdir,
name=wandb_name,
id=id,
resume="allow",
)
else:
wandb.init(
project=wandb_project,
sync_tensorboard=True,
name=wandb_name,
dir=logdir,
)
writer = SummaryWriter(Path(logdir) / "tensorboard")
return writer