gogo2/NN/start_tensorboard.py
Dobromir Popov 0042581275 new nn wip
2025-03-25 13:38:25 +02:00

88 lines
2.6 KiB
Python

#!/usr/bin/env python
"""
Start TensorBoard for monitoring neural network training
"""
import os
import sys
import subprocess
import webbrowser
from time import sleep
def start_tensorboard(logdir="NN/models/saved/logs", port=6006, open_browser=True):
"""
Start TensorBoard in a subprocess
Args:
logdir: Directory containing TensorBoard logs
port: Port to run TensorBoard on
open_browser: Whether to open a browser automatically
"""
# Make sure the log directory exists
os.makedirs(logdir, exist_ok=True)
# Create command
cmd = [
sys.executable,
"-m",
"tensorboard.main",
f"--logdir={logdir}",
f"--port={port}",
"--bind_all"
]
print(f"Starting TensorBoard with logs from {logdir} on port {port}")
print(f"Command: {' '.join(cmd)}")
# Start TensorBoard in a subprocess
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
universal_newlines=True
)
# Wait for TensorBoard to start up
for line in process.stdout:
print(line.strip())
if "TensorBoard" in line and "http://" in line:
# TensorBoard is running, extract the URL
url = None
for part in line.split():
if part.startswith(("http://", "https://")):
url = part
break
# Open browser if requested and URL found
if open_browser and url:
print(f"Opening TensorBoard in browser: {url}")
webbrowser.open(url)
break
# Return the process for the caller to manage
return process
if __name__ == "__main__":
import argparse
# Parse command line arguments
parser = argparse.ArgumentParser(description="Start TensorBoard for NN training visualization")
parser.add_argument("--logdir", default="NN/models/saved/logs", help="Directory containing TensorBoard logs")
parser.add_argument("--port", type=int, default=6006, help="Port to run TensorBoard on")
parser.add_argument("--no-browser", action="store_true", help="Don't open browser automatically")
args = parser.parse_args()
# Start TensorBoard
process = start_tensorboard(args.logdir, args.port, not args.no_browser)
try:
# Keep the script running until Ctrl+C
print("TensorBoard is running. Press Ctrl+C to stop.")
while True:
sleep(1)
except KeyboardInterrupt:
print("Stopping TensorBoard...")
process.terminate()
process.wait()