trnbl.loggers.tensorboard API documentation

docs for trnbl v0.1.0
View Source on GitHub

# trnbl.loggers.tensorboard

View Source
View on GitHub

      1import datetime
      2from typing import Any
      3from pathlib import Path
      4import json
      5import hashlib
      6
      7from torch.utils.tensorboard import SummaryWriter
      8
      9from trnbl.loggers.base import TrainingLoggerBase, rand_syllabic_string
     10
     11
     12class TensorBoardLogger(TrainingLoggerBase):
     13	def __init__(
     14		self,
     15		log_dir: str | Path,
     16		train_config: dict | None = None,
     17		name: str | None = None,
     18		**kwargs,
     19	) -> None:
     20		train_config_json: str = json.dumps(train_config, indent="\t")
     21
     22		if name is None:
     23			_run_hash: str = hashlib.md5(train_config_json.encode()).hexdigest()
     24			name = f"h{_run_hash[:5]}-{datetime.datetime.now().strftime('%y%m%d_%H%M')}-{rand_syllabic_string()}"
     25
     26		log_dir = (Path(log_dir) / name).as_posix()
     27
     28		# Initialize the TensorBoard SummaryWriter with the specified log directory
     29		self._writer: SummaryWriter = SummaryWriter(log_dir=log_dir, **kwargs)
     30
     31		# Store the run path
     32		self._run_path: Path = Path(log_dir)
     33
     34		# Initialize the global step counter
     35		self._global_step: int = 0
     36
     37		# Log the training configuration
     38		self._self_writer_add_text("config", train_config_json)
     39		self._self_writer_add_text("name", name)
     40		with open(self._run_path / "config.json", "w") as f:
     41			f.write(train_config_json)
     42
     43	def _self_writer_add_text(self, tag: str, message: str, **kwargs) -> None:
     44		self._writer.add_text(
     45			tag,
     46			message + ("" if not kwargs else "\n" + json.dumps(kwargs, indent=4)),
     47			global_step=self._global_step,
     48		)
     49
     50	def debug(self, message: str, **kwargs) -> None:
     51		self._self_writer_add_text("debug", message, **kwargs)
     52
     53	def message(self, message: str, **kwargs) -> None:
     54		self._self_writer_add_text("message", message, **kwargs)
     55
     56		# Also print the message
     57		print(message)
     58
     59	def metrics(self, data: dict[str, Any]) -> None:
     60		# Log a dictionary of metrics using add_scalar in TensorBoard
     61		for key, value in data.items():
     62			self._writer.add_scalar(key, value, global_step=self._global_step)
     63
     64		# Increment the global step counter
     65		self._global_step += 1
     66
     67	def artifact(
     68		self,
     69		path: Path,
     70		type: str,
     71		aliases: list[str] | None = None,
     72		metadata: dict | None = None,
     73	) -> None:
     74		# Log an artifact file using add_artifact in TensorBoard
     75		self._writer.add_text(
     76			tag="artifact",
     77			text_string=json.dumps(
     78				dict(
     79					timestamp=datetime.datetime.now().isoformat(),
     80					path=path.as_posix(),
     81					type=type,
     82					aliases=aliases,
     83					metadata=metadata if metadata else {},
     84				)
     85			),
     86			global_step=self._global_step,
     87		)
     88
     89	@property
     90	def url(self) -> str:
     91		# Return the command to launch TensorBoard with the specified log directory
     92		return f"tensorboard --logdir={self._run_path}"
     93
     94	@property
     95	def run_path(self) -> Path:
     96		# Return the run path
     97		return self._run_path
     98
     99	def flush(self) -> None:
    100		self._writer.flush()
    101
    102	def finish(self) -> None:
    103		# Flush and close the TensorBoard SummaryWriter
    104		self._writer.flush()
    105		self._writer.close()

class TensorBoardLogger(trnbl.loggers.base.TrainingLoggerBase):
View Source
View on GitHub

     13class TensorBoardLogger(TrainingLoggerBase):
     14	def __init__(
     15		self,
     16		log_dir: str | Path,
     17		train_config: dict | None = None,
     18		name: str | None = None,
     19		**kwargs,
     20	) -> None:
     21		train_config_json: str = json.dumps(train_config, indent="\t")
     22
     23		if name is None:
     24			_run_hash: str = hashlib.md5(train_config_json.encode()).hexdigest()
     25			name = f"h{_run_hash[:5]}-{datetime.datetime.now().strftime('%y%m%d_%H%M')}-{rand_syllabic_string()}"
     26
     27		log_dir = (Path(log_dir) / name).as_posix()
     28
     29		# Initialize the TensorBoard SummaryWriter with the specified log directory
     30		self._writer: SummaryWriter = SummaryWriter(log_dir=log_dir, **kwargs)
     31
     32		# Store the run path
     33		self._run_path: Path = Path(log_dir)
     34
     35		# Initialize the global step counter
     36		self._global_step: int = 0
     37
     38		# Log the training configuration
     39		self._self_writer_add_text("config", train_config_json)
     40		self._self_writer_add_text("name", name)
     41		with open(self._run_path / "config.json", "w") as f:
     42			f.write(train_config_json)
     43
     44	def _self_writer_add_text(self, tag: str, message: str, **kwargs) -> None:
     45		self._writer.add_text(
     46			tag,
     47			message + ("" if not kwargs else "\n" + json.dumps(kwargs, indent=4)),
     48			global_step=self._global_step,
     49		)
     50
     51	def debug(self, message: str, **kwargs) -> None:
     52		self._self_writer_add_text("debug", message, **kwargs)
     53
     54	def message(self, message: str, **kwargs) -> None:
     55		self._self_writer_add_text("message", message, **kwargs)
     56
     57		# Also print the message
     58		print(message)
     59
     60	def metrics(self, data: dict[str, Any]) -> None:
     61		# Log a dictionary of metrics using add_scalar in TensorBoard
     62		for key, value in data.items():
     63			self._writer.add_scalar(key, value, global_step=self._global_step)
     64
     65		# Increment the global step counter
     66		self._global_step += 1
     67
     68	def artifact(
     69		self,
     70		path: Path,
     71		type: str,
     72		aliases: list[str] | None = None,
     73		metadata: dict | None = None,
     74	) -> None:
     75		# Log an artifact file using add_artifact in TensorBoard
     76		self._writer.add_text(
     77			tag="artifact",
     78			text_string=json.dumps(
     79				dict(
     80					timestamp=datetime.datetime.now().isoformat(),
     81					path=path.as_posix(),
     82					type=type,
     83					aliases=aliases,
     84					metadata=metadata if metadata else {},
     85				)
     86			),
     87			global_step=self._global_step,
     88		)
     89
     90	@property
     91	def url(self) -> str:
     92		# Return the command to launch TensorBoard with the specified log directory
     93		return f"tensorboard --logdir={self._run_path}"
     94
     95	@property
     96	def run_path(self) -> Path:
     97		# Return the run path
     98		return self._run_path
     99
    100	def flush(self) -> None:
    101		self._writer.flush()
    102
    103	def finish(self) -> None:
    104		# Flush and close the TensorBoard SummaryWriter
    105		self._writer.flush()
    106		self._writer.close()

Base class for training loggers

TensorBoardLogger( log_dir: str | pathlib.Path, train_config: dict | None = None, name: str | None = None, **kwargs)
View Source
View on GitHub

    14	def __init__(
    15		self,
    16		log_dir: str | Path,
    17		train_config: dict | None = None,
    18		name: str | None = None,
    19		**kwargs,
    20	) -> None:
    21		train_config_json: str = json.dumps(train_config, indent="\t")
    22
    23		if name is None:
    24			_run_hash: str = hashlib.md5(train_config_json.encode()).hexdigest()
    25			name = f"h{_run_hash[:5]}-{datetime.datetime.now().strftime('%y%m%d_%H%M')}-{rand_syllabic_string()}"
    26
    27		log_dir = (Path(log_dir) / name).as_posix()
    28
    29		# Initialize the TensorBoard SummaryWriter with the specified log directory
    30		self._writer: SummaryWriter = SummaryWriter(log_dir=log_dir, **kwargs)
    31
    32		# Store the run path
    33		self._run_path: Path = Path(log_dir)
    34
    35		# Initialize the global step counter
    36		self._global_step: int = 0
    37
    38		# Log the training configuration
    39		self._self_writer_add_text("config", train_config_json)
    40		self._self_writer_add_text("name", name)
    41		with open(self._run_path / "config.json", "w") as f:
    42			f.write(train_config_json)

def debug(self, message: str, **kwargs) -> None:
View Source
View on GitHub

    51	def debug(self, message: str, **kwargs) -> None:
    52		self._self_writer_add_text("debug", message, **kwargs)

log a debug message which will be saved, but not printed

def message(self, message: str, **kwargs) -> None:
View Source
View on GitHub

    54	def message(self, message: str, **kwargs) -> None:
    55		self._self_writer_add_text("message", message, **kwargs)
    56
    57		# Also print the message
    58		print(message)

log a progress message, which will be printed to stdout

def metrics(self, data: dict[str, typing.Any]) -> None:
View Source
View on GitHub

    60	def metrics(self, data: dict[str, Any]) -> None:
    61		# Log a dictionary of metrics using add_scalar in TensorBoard
    62		for key, value in data.items():
    63			self._writer.add_scalar(key, value, global_step=self._global_step)
    64
    65		# Increment the global step counter
    66		self._global_step += 1

Log a dictionary of metrics

def artifact( self, path: pathlib.Path, type: str, aliases: list[str] | None = None, metadata: dict | None = None) -> None:
View Source
View on GitHub

    68	def artifact(
    69		self,
    70		path: Path,
    71		type: str,
    72		aliases: list[str] | None = None,
    73		metadata: dict | None = None,
    74	) -> None:
    75		# Log an artifact file using add_artifact in TensorBoard
    76		self._writer.add_text(
    77			tag="artifact",
    78			text_string=json.dumps(
    79				dict(
    80					timestamp=datetime.datetime.now().isoformat(),
    81					path=path.as_posix(),
    82					type=type,
    83					aliases=aliases,
    84					metadata=metadata if metadata else {},
    85				)
    86			),
    87			global_step=self._global_step,
    88		)

log an artifact from a file

url: str
View Source
View on GitHub

    90	@property
    91	def url(self) -> str:
    92		# Return the command to launch TensorBoard with the specified log directory
    93		return f"tensorboard --logdir={self._run_path}"

Get the URL for the current logging run

run_path: pathlib.Path
View Source
View on GitHub

    95	@property
    96	def run_path(self) -> Path:
    97		# Return the run path
    98		return self._run_path

Get the path to the current logging run

def flush(self) -> None:
View Source
View on GitHub

    100	def flush(self) -> None:
    101		self._writer.flush()

Flush the logger

def finish(self) -> None:
View Source
View on GitHub

    103	def finish(self) -> None:
    104		# Flush and close the TensorBoard SummaryWriter
    105		self._writer.flush()
    106		self._writer.close()

Finish logging

##### Inherited Members

trnbl.loggers.base.TrainingLoggerBase
  warning
  error
  get_mem_usage
  spinner_task