vscode-python-manager/pythonFiles/tensorboard_launcher.py at main · codegzx/vscode-python-manager · GitHub

36 lines (29 loc) · 991 Bytes

import time

import sys

import os

import mimetypes

from tensorboard import program

def main(logdir):

# Environment variable for PyTorch profiler TensorBoard plugin

# to detect when it's running inside VS Code

os.environ["VSCODE_TENSORBOARD_LAUNCH"] = "1"

# Work around incorrectly configured MIME types on Windows

mimetypes.add_type("application/javascript", ".js")

# Start TensorBoard using their Python API

tb = program.TensorBoard()

tb.configure(bind_all=False, logdir=logdir)

url = tb.launch()

sys.stdout.write("TensorBoard started at %s\n" % (url))

sys.stdout.flush()

while True:

try:

time.sleep(60)

except KeyboardInterrupt:

break

sys.stdout.write("TensorBoard is shutting down")

sys.stdout.flush()

if __name__ == "__main__":

if len(sys.argv) == 2:

logdir = str(sys.argv[1])

sys.stdout.write("Starting TensorBoard with logdir %s" % (logdir))

main(logdir)

Expand file tree