|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import logging |
|
|
import os |
|
|
import psutil |
|
|
import time |
|
|
import subprocess |
|
|
from torch.utils.tensorboard import SummaryWriter |
|
|
|
|
|
DEFAULT_TENSORBOARD_PORT = 6006 |
|
|
|
|
|
|
|
|
|
|
|
def get_writer(path: str, tensorboard_port: int | bool, logger: logging.Logger = None): |
|
|
""" |
|
|
Sets up a TensorBoard logging and checkpoint directory for PyTorch. |
|
|
|
|
|
This function clears the specified directory, creates subdirectories for TensorBoard logs |
|
|
and model checkpoints, ensuring a clean environment for running new training sessions. |
|
|
|
|
|
Args: |
|
|
path (str): The root directory where TensorBoard logs and checkpoints will be stored. |
|
|
tensorboard_port (int): The port on which to run the TensorBoard. |
|
|
logger (logging.Logger): The logger that traces the logging information. |
|
|
|
|
|
Returns: |
|
|
tuple: A tuple containing the TensorBoard SummaryWriter object and the path for checkpoints. |
|
|
|
|
|
Example: |
|
|
>>> tensor_writer, checkpoint_dir = get_writer('/path/to/tensorboard/') |
|
|
""" |
|
|
|
|
|
if tensorboard_port is True: |
|
|
tensorboard_port = DEFAULT_TENSORBOARD_PORT |
|
|
elif tensorboard_port is False: |
|
|
return None, os.path.join(path, 'checkpoints') |
|
|
|
|
|
|
|
|
logs_path = os.path.join(path, 'logs') |
|
|
checkpoints_path = os.path.join(path, 'checkpoints') |
|
|
os.makedirs(logs_path, exist_ok=True) |
|
|
os.makedirs(checkpoints_path, exist_ok=True) |
|
|
|
|
|
|
|
|
writer = SummaryWriter(log_dir=logs_path) |
|
|
|
|
|
|
|
|
if logger is not None: |
|
|
logger.info(f"TensorBoard logs will be stored in: {logs_path}") |
|
|
logger.info(f"Model checkpoints will be stored in: {checkpoints_path}") |
|
|
|
|
|
|
|
|
for conn in psutil.net_connections(kind='inet'): |
|
|
if conn.laddr.port == tensorboard_port and conn.status == psutil.CONN_LISTEN: |
|
|
if logger is not None: |
|
|
logger.warning(f"Killing already running TensorBoard process with PID {conn.pid}") |
|
|
p = psutil.Process(conn.pid) |
|
|
p.terminate() |
|
|
p.wait(timeout=3) |
|
|
time.sleep(5) |
|
|
process = subprocess.Popen(f'tensorboard --logdir={logs_path} --host=0.0.0.0 --port={tensorboard_port}', |
|
|
shell=True) |
|
|
if logger is not None: |
|
|
logger.info(f'TensorBoard running at http://0.0.0.0:{tensorboard_port}/ (pid={process.pid})') |
|
|
|
|
|
return writer, checkpoints_path |
|
|
|
|
|
|
|
|
|
|
|
|
