from __future__ import annotations import io import re import time from typing import TYPE_CHECKING, Any import wandb import wandb.util from wandb.sdk.lib import telemetry if TYPE_CHECKING: import numpy as np from wandb.sdk.internal.tb_watcher import TBHistory # We have at least the default namestep and a global step to track # TODO: reset this structure on wandb.finish STEPS: dict[str, dict[str, Any]] = { "": {"step": 0}, "global": {"step": 0, "last_log": None}, } # TODO(cling): Set these when tensorboard behavior is configured. # We support rate limited logging by setting this to number of seconds, # can be a floating point. RATE_LIMIT_SECONDS: float | int | None = None IGNORE_KINDS = ["graphs"] tensor_util = wandb.util.get_module("tensorboard.util.tensor_util") # prefer tensorboard, fallback to protobuf in tensorflow when tboard isn't available pb = wandb.util.get_module( "tensorboard.compat.proto.summary_pb2" ) or wandb.util.get_module("tensorflow.core.framework.summary_pb2") Summary = pb.Summary if pb else None def make_ndarray(tensor: Any) -> np.ndarray | None: if tensor_util: res = tensor_util.make_ndarray(tensor) # Tensorboard can log generic objects, and we don't want to save them if res.dtype == "object": return None else: return res # type: ignore else: wandb.termwarn( "Can't convert tensor summary, upgrade tensorboard with `pip" " install tensorboard --upgrade`" ) return None def namespaced_tag(tag: str, namespace: str = "") -> str: if not namespace: return tag else: return namespace + "/" + tag def history_image_key(key: str, namespace: str = "") -> str: """Convert invalid filesystem characters to _ for use in History keys. Unfortunately this means currently certain image keys will collide silently. We implement this mapping up here in the TensorFlow stuff rather than in the History stuff so that we don't have to store a mapping anywhere from the original keys to the safe ones. """ return namespaced_tag(re.sub(r"[/\\]", "_", key), namespace) def tf_summary_to_dict( # noqa: C901 tf_summary_str_or_pb: Any, namespace: str = "" ) -> dict[str, Any] | None: """Convert a Tensorboard Summary to a dictionary. Accepts a tensorflow.summary.Summary, one encoded as a string, or a list of such encoded as strings. """ values = {} if hasattr(tf_summary_str_or_pb, "summary"): summary_pb = tf_summary_str_or_pb.summary values[namespaced_tag("global_step", namespace)] = tf_summary_str_or_pb.step values["_timestamp"] = tf_summary_str_or_pb.wall_time elif isinstance(tf_summary_str_or_pb, (str, bytes, bytearray)): summary_pb = Summary() summary_pb.ParseFromString(tf_summary_str_or_pb) elif hasattr(tf_summary_str_or_pb, "__iter__"): summary_pb = [Summary() for _ in range(len(tf_summary_str_or_pb))] for i, summary in enumerate(tf_summary_str_or_pb): summary_pb[i].ParseFromString(summary) if i > 0: summary_pb[0].MergeFrom(summary_pb[i]) summary_pb = summary_pb[0] else: summary_pb = tf_summary_str_or_pb if not hasattr(summary_pb, "value") or len(summary_pb.value) == 0: # Ignore these, caller is responsible for handling None return None def encode_images(_img_strs: list[bytes], _value: Any) -> None: try: from PIL import Image except ImportError: wandb.termwarn( "Install pillow if you are logging images with Tensorboard. " "To install, run `pip install pillow`.", repeat=False, ) return None if len(_img_strs) == 0: return None images: list[wandb.Video | wandb.Image] = [] for _img_str in _img_strs: # Supports gifs from TensorboardX if _img_str.startswith(b"GIF"): images.append(wandb.Video(io.BytesIO(_img_str), format="gif")) else: images.append(wandb.Image(Image.open(io.BytesIO(_img_str)))) tag_idx = _value.tag.rsplit("/", 1) if len(tag_idx) > 1 and tag_idx[1].isdigit(): tag, idx = tag_idx values.setdefault(history_image_key(tag, namespace), []).extend(images) else: values[history_image_key(_value.tag, namespace)] = images return None for value in summary_pb.value: kind = value.WhichOneof("value") if kind in IGNORE_KINDS: continue if kind == "simple_value": values[namespaced_tag(value.tag, namespace)] = value.simple_value elif kind == "tensor": plugin_name = value.metadata.plugin_data.plugin_name if plugin_name == "scalars" or plugin_name == "": values[namespaced_tag(value.tag, namespace)] = make_ndarray( value.tensor ) elif plugin_name == "images": img_strs = value.tensor.string_val[2:] # First two items are dims. encode_images(img_strs, value) elif plugin_name == "histograms": # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/histogram/summary_v2.py#L15-L26 ndarray = make_ndarray(value.tensor) if ndarray is None: continue shape = ndarray.shape counts = [] bins = [] if shape[0] > 1: bins.append(ndarray[0][0]) # Add the left most edge for v in ndarray: counts.append(v[2]) bins.append(v[1]) # Add the right most edges elif shape[0] == 1: counts = [ndarray[0][2]] bins = ndarray[0][:2] if len(counts) > 0: try: # TODO: we should just re-bin if there are too many buckets values[namespaced_tag(value.tag, namespace)] = wandb.Histogram( np_histogram=(counts, bins) # type: ignore ) except ValueError: wandb.termwarn( f'Not logging key "{namespaced_tag(value.tag, namespace)}". ' f"Histograms must have fewer than {wandb.Histogram.MAX_LENGTH} bins", repeat=False, ) elif plugin_name == "pr_curves": pr_curve_data = make_ndarray(value.tensor) if pr_curve_data is None: continue precision = pr_curve_data[-2, :].tolist() recall = pr_curve_data[-1, :].tolist() # TODO: (kdg) implement spec for showing additional info in tool tips # true_pos = pr_curve_data[1,:] # false_pos = pr_curve_data[2,:] # true_neg = pr_curve_data[1,:] # false_neg = pr_curve_data[1,:] # threshold = [1.0 / n for n in range(len(true_pos), 0, -1)] # min of each in case tensorboard ever changes their pr_curve # to allow for different length outputs data = [] for i in range(min(len(precision), len(recall))): # drop additional threshold values if they exist if precision[i] != 0 or recall[i] != 0: data.append((recall[i], precision[i])) # sort data so custom chart looks the same as tb generated pr curve # ascending recall, descending precision for the same recall values data = sorted(data, key=lambda x: (x[0], -x[1])) data_table = wandb.Table(data=data, columns=["recall", "precision"]) name = namespaced_tag(value.tag, namespace) values[name] = wandb.plot_table( "wandb/line/v0", data_table, {"x": "recall", "y": "precision"}, {"title": f"{name} Precision v. Recall"}, ) elif kind == "image": img_str = value.image.encoded_image_string encode_images([img_str], value) # Coming soon... # elif kind == "audio": # audio = wandb.Audio( # six.BytesIO(value.audio.encoded_audio_string), # sample_rate=value.audio.sample_rate, # content_type=value.audio.content_type, # ) elif kind == "histo": tag = namespaced_tag(value.tag, namespace) if len(value.histo.bucket_limit) >= 3: first = ( value.histo.bucket_limit[0] + value.histo.bucket_limit[0] - value.histo.bucket_limit[1] ) last = ( value.histo.bucket_limit[-2] + value.histo.bucket_limit[-2] - value.histo.bucket_limit[-3] ) np_histogram = ( list(value.histo.bucket), [first] + value.histo.bucket_limit[:-1] + [last], ) try: # TODO: we should just re-bin if there are too many buckets values[tag] = wandb.Histogram(np_histogram=np_histogram) # type: ignore except ValueError: wandb.termwarn( f"Not logging key {tag!r}. " f"Histograms must have fewer than {wandb.Histogram.MAX_LENGTH} bins", repeat=False, ) else: # TODO: is there a case where we can render this? wandb.termwarn( f"Not logging key {tag!r}. Found a histogram with only 2 bins.", repeat=False, ) # TODO(jhr): figure out how to share this between userspace and internal process or dont # elif value.tag == "_hparams_/session_start_info": # if wandb.util.get_module("tensorboard.plugins.hparams"): # from tensorboard.plugins.hparams import plugin_data_pb2 # # plugin_data = plugin_data_pb2.HParamsPluginData() # # plugin_data.ParseFromString(value.metadata.plugin_data.content) # for key, param in six.iteritems(plugin_data.session_start_info.hparams): # if not wandb.run.config.get(key): # wandb.run.config[key] = ( # param.number_value or param.string_value or param.bool_value # ) # else: # wandb.termerror( # "Received hparams tf.summary, but could not import " # "the hparams plugin from tensorboard" # ) return values def reset_state() -> None: """Internal method for resetting state, called by wandb.finish().""" global STEPS STEPS = {"": {"step": 0}, "global": {"step": 0, "last_log": None}} def _log( tf_summary_str_or_pb: Any, history: TBHistory | None = None, step: int = 0, namespace: str = "", **kwargs: Any, ) -> None: """Logs a tfsummary to wandb. Can accept a tf summary string or parsed event. Will use wandb.run.history unless a history object is passed. Can optionally namespace events. Results are committed when step increases for this namespace. NOTE: This assumes that events being passed in are in chronological order """ global STEPS global RATE_LIMIT_SECONDS # To handle multiple global_steps, we keep track of them here instead # of the global log last_step = STEPS.get(namespace, {"step": 0}) # Commit our existing data if this namespace increased its step commit = False if last_step["step"] < step: commit = True log_dict = tf_summary_to_dict(tf_summary_str_or_pb, namespace) if log_dict is None: # not an event, just return return # Pass timestamp to history for loading historic data timestamp = log_dict.get("_timestamp", time.time()) # Store our initial timestamp if STEPS["global"]["last_log"] is None: STEPS["global"]["last_log"] = timestamp # Rollup events that share the same step across namespaces if commit and step == STEPS["global"]["step"]: commit = False # Always add the biggest global_step key for non-default namespaces if step > STEPS["global"]["step"]: STEPS["global"]["step"] = step if namespace != "": log_dict["global_step"] = STEPS["global"]["step"] # Keep internal step counter STEPS[namespace] = {"step": step} if commit: # Only commit our data if we're below the rate limit or don't have one if ( RATE_LIMIT_SECONDS is None or timestamp - STEPS["global"]["last_log"] >= RATE_LIMIT_SECONDS ): if history is None: if wandb.run is not None: wandb.run._log({}) else: history.add({}) STEPS["global"]["last_log"] = timestamp if history is None: if wandb.run is not None: wandb.run._log(log_dict, commit=False) else: history._row_update(log_dict) def log(tf_summary_str_or_pb: Any, step: int = 0, namespace: str = "") -> None: if wandb.run is None: raise wandb.Error( "You must call `wandb.init()` before calling `wandb.tensorflow.log`" ) with telemetry.context() as tel: tel.feature.tensorboard_log = True _log(tf_summary_str_or_pb, namespace=namespace, step=step)