wandb 0.18.0rc1__py3-none-any.whl → 0.18.2__py3-none-any.whl

wandb/sdk/data_types/graph.py

@@ -0,0 +1,405 @@
+import codecs
+import os
+import pprint
+
+from wandb import util
+from wandb.sdk.data_types._private import MEDIA_TMP
+from wandb.sdk.data_types.base_types.media import Media, _numpy_arrays_to_lists
+from wandb.sdk.data_types.base_types.wb_value import WBValue
+from wandb.sdk.lib import runid
+
+
+def _nest(thing):
+    # Use tensorflows nest function if available, otherwise just wrap object in an array"""
+
+    tfutil = util.get_module("tensorflow.python.util")
+    if tfutil:
+        return tfutil.nest.flatten(thing)
+    else:
+        return [thing]
+
+
+class Edge(WBValue):
+    """Edge used in `Graph`."""
+
+    def __init__(self, from_node, to_node):
+        self._attributes = {}
+        self.from_node = from_node
+        self.to_node = to_node
+
+    def __repr__(self):
+        temp_attr = dict(self._attributes)
+        del temp_attr["from_node"]
+        del temp_attr["to_node"]
+        temp_attr["from_id"] = self.from_node.id
+        temp_attr["to_id"] = self.to_node.id
+        return str(temp_attr)
+
+    def to_json(self, run=None):
+        return [self.from_node.id, self.to_node.id]
+
+    @property
+    def name(self):
+        """Optional, not necessarily unique."""
+        return self._attributes.get("name")
+
+    @name.setter
+    def name(self, val):
+        self._attributes["name"] = val
+        return val
+
+    @property
+    def from_node(self):
+        return self._attributes.get("from_node")
+
+    @from_node.setter
+    def from_node(self, val):
+        self._attributes["from_node"] = val
+        return val
+
+    @property
+    def to_node(self):
+        return self._attributes.get("to_node")
+
+    @to_node.setter
+    def to_node(self, val):
+        self._attributes["to_node"] = val
+        return val
+
+
+class Node(WBValue):
+    """Node used in `Graph`."""
+
+    def __init__(
+        self,
+        id=None,
+        name=None,
+        class_name=None,
+        size=None,
+        parameters=None,
+        output_shape=None,
+        is_output=None,
+        num_parameters=None,
+        node=None,
+    ):
+        self._attributes = {"name": None}
+        self.in_edges = {}  # indexed by source node id
+        self.out_edges = {}  # indexed by dest node id
+        # optional object (e.g. PyTorch Parameter or Module) that this Node represents
+        self.obj = None
+
+        if node is not None:
+            self._attributes.update(node._attributes)
+            del self._attributes["id"]
+            self.obj = node.obj
+
+        if id is not None:
+            self.id = id
+        if name is not None:
+            self.name = name
+        if class_name is not None:
+            self.class_name = class_name
+        if size is not None:
+            self.size = size
+        if parameters is not None:
+            self.parameters = parameters
+        if output_shape is not None:
+            self.output_shape = output_shape
+        if is_output is not None:
+            self.is_output = is_output
+        if num_parameters is not None:
+            self.num_parameters = num_parameters
+
+    def to_json(self, run=None):
+        return self._attributes
+
+    def __repr__(self):
+        return repr(self._attributes)
+
+    @property
+    def id(self):
+        """Must be unique in the graph."""
+        return self._attributes.get("id")
+
+    @id.setter
+    def id(self, val):
+        self._attributes["id"] = val
+        return val
+
+    @property
+    def name(self):
+        """Usually the type of layer or sublayer."""
+        return self._attributes.get("name")
+
+    @name.setter
+    def name(self, val):
+        self._attributes["name"] = val
+        return val
+
+    @property
+    def class_name(self):
+        """Usually the type of layer or sublayer."""
+        return self._attributes.get("class_name")
+
+    @class_name.setter
+    def class_name(self, val):
+        self._attributes["class_name"] = val
+        return val
+
+    @property
+    def functions(self):
+        return self._attributes.get("functions", [])
+
+    @functions.setter
+    def functions(self, val):
+        self._attributes["functions"] = val
+        return val
+
+    @property
+    def parameters(self):
+        return self._attributes.get("parameters", [])
+
+    @parameters.setter
+    def parameters(self, val):
+        self._attributes["parameters"] = val
+        return val
+
+    @property
+    def size(self):
+        return self._attributes.get("size")
+
+    @size.setter
+    def size(self, val):
+        """Tensor size."""
+        self._attributes["size"] = tuple(val)
+        return val
+
+    @property
+    def output_shape(self):
+        return self._attributes.get("output_shape")
+
+    @output_shape.setter
+    def output_shape(self, val):
+        """Tensor output_shape."""
+        self._attributes["output_shape"] = val
+        return val
+
+    @property
+    def is_output(self):
+        return self._attributes.get("is_output")
+
+    @is_output.setter
+    def is_output(self, val):
+        """Tensor is_output."""
+        self._attributes["is_output"] = val
+        return val
+
+    @property
+    def num_parameters(self):
+        return self._attributes.get("num_parameters")
+
+    @num_parameters.setter
+    def num_parameters(self, val):
+        """Tensor num_parameters."""
+        self._attributes["num_parameters"] = val
+        return val
+
+    @property
+    def child_parameters(self):
+        return self._attributes.get("child_parameters")
+
+    @child_parameters.setter
+    def child_parameters(self, val):
+        """Tensor child_parameters."""
+        self._attributes["child_parameters"] = val
+        return val
+
+    @property
+    def is_constant(self):
+        return self._attributes.get("is_constant")
+
+    @is_constant.setter
+    def is_constant(self, val):
+        """Tensor is_constant."""
+        self._attributes["is_constant"] = val
+        return val
+
+    @classmethod
+    def from_keras(cls, layer):
+        node = cls()
+
+        try:
+            output_shape = layer.output_shape
+        except AttributeError:
+            output_shape = ["multiple"]
+
+        node.id = layer.name
+        node.name = layer.name
+        node.class_name = layer.__class__.__name__
+        node.output_shape = output_shape
+        node.num_parameters = layer.count_params()
+
+        return node
+
+
+class Graph(Media):
+    """Wandb class for graphs.
+
+    This class is typically used for saving and displaying neural net models.  It
+    represents the graph as an array of nodes and edges.  The nodes can have
+    labels that can be visualized by wandb.
+
+    Examples:
+        Import a keras model:
+        ```
+        Graph.from_keras(keras_model)
+        ```
+
+    Attributes:
+        format (string): Format to help wandb display the graph nicely.
+        nodes ([wandb.Node]): List of wandb.Nodes
+        nodes_by_id (dict): dict of ids -> nodes
+        edges ([(wandb.Node, wandb.Node)]): List of pairs of nodes interpreted as edges
+        loaded (boolean): Flag to tell whether the graph is completely loaded
+        root (wandb.Node): root node of the graph
+    """
+
+    _log_type = "graph-file"
+
+    def __init__(self, format="keras"):
+        super().__init__()
+        # LB: TODO: I think we should factor criterion and criterion_passed out
+        self.format = format
+        self.nodes = []
+        self.nodes_by_id = {}
+        self.edges = []
+        self.loaded = False
+        self.criterion = None
+        self.criterion_passed = False
+        self.root = None  # optional root Node if applicable
+
+    def _to_graph_json(self, run=None):
+        # Needs to be its own function for tests
+        return {
+            "format": self.format,
+            "nodes": [node.to_json() for node in self.nodes],
+            "edges": [edge.to_json() for edge in self.edges],
+        }
+
+    def bind_to_run(self, *args, **kwargs):
+        data = self._to_graph_json()
+        tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".graph.json")
+        data = _numpy_arrays_to_lists(data)
+        with codecs.open(tmp_path, "w", encoding="utf-8") as fp:
+            util.json_dump_safer(data, fp)
+        self._set_file(tmp_path, is_tmp=True, extension=".graph.json")
+        if self.is_bound():
+            return
+        super().bind_to_run(*args, **kwargs)
+
+    @classmethod
+    def get_media_subdir(cls):
+        return os.path.join("media", "graph")
+
+    def to_json(self, run):
+        json_dict = super().to_json(run)
+        json_dict["_type"] = self._log_type
+        return json_dict
+
+    def __getitem__(self, nid):
+        return self.nodes_by_id[nid]
+
+    def pprint(self):
+        for edge in self.edges:
+            pprint.pprint(edge.attributes)
+        for node in self.nodes:
+            pprint.pprint(node.attributes)
+
+    def add_node(self, node=None, **node_kwargs):
+        if node is None:
+            node = Node(**node_kwargs)
+        elif node_kwargs:
+            raise ValueError(
+                f"Only pass one of either node ({node}) or other keyword arguments ({node_kwargs})"
+            )
+        self.nodes.append(node)
+        self.nodes_by_id[node.id] = node
+
+        return node
+
+    def add_edge(self, from_node, to_node):
+        edge = Edge(from_node, to_node)
+        self.edges.append(edge)
+
+        return edge
+
+    @classmethod
+    def from_keras(cls, model):
+        # TODO: his method requires a refactor to work with the keras 3.
+        graph = cls()
+        # Shamelessly copied (then modified) from keras/keras/utils/layer_utils.py
+        sequential_like = cls._is_sequential(model)
+
+        relevant_nodes = None
+        if not sequential_like:
+            relevant_nodes = []
+            for v in model._nodes_by_depth.values():
+                relevant_nodes += v
+
+        layers = model.layers
+        for i in range(len(layers)):
+            node = Node.from_keras(layers[i])
+            if hasattr(layers[i], "_inbound_nodes"):
+                for in_node in layers[i]._inbound_nodes:
+                    if relevant_nodes and in_node not in relevant_nodes:
+                        # node is not part of the current network
+                        continue
+                    for in_layer in _nest(in_node.inbound_layers):
+                        inbound_keras_node = Node.from_keras(in_layer)
+
+                        if inbound_keras_node.id not in graph.nodes_by_id:
+                            graph.add_node(inbound_keras_node)
+                        inbound_node = graph.nodes_by_id[inbound_keras_node.id]
+
+                        graph.add_edge(inbound_node, node)
+            graph.add_node(node)
+        return graph
+
+    @classmethod
+    def _is_sequential(cls, model):
+        sequential_like = True
+
+        if (
+            model.__class__.__name__ != "Sequential"
+            and hasattr(model, "_is_graph_network")
+            and model._is_graph_network
+        ):
+            nodes_by_depth = model._nodes_by_depth.values()
+            nodes = []
+            for v in nodes_by_depth:
+                # TensorFlow2 doesn't insure inbound is always a list
+                inbound = v[0].inbound_layers
+                if not hasattr(inbound, "__len__"):
+                    inbound = [inbound]
+                if (len(v) > 1) or (len(v) == 1 and len(inbound) > 1):
+                    # if the model has multiple nodes
+                    # or if the nodes have multiple inbound_layers
+                    # the model is no longer sequential
+                    sequential_like = False
+                    break
+                nodes += v
+            if sequential_like:
+                # search for shared layers
+                for layer in model.layers:
+                    flag = False
+                    if hasattr(layer, "_inbound_nodes"):
+                        for node in layer._inbound_nodes:
+                            if node in nodes:
+                                if flag:
+                                    sequential_like = False
+                                    break
+                                else:
+                                    flag = True
+                    if not sequential_like:
+                        break
+        return sequential_like

wandb/sdk/data_types/image.py

@@ -10,6 +10,7 @@ from wandb import util
 from wandb.sdk.lib import hashutil, runid
 from wandb.sdk.lib.paths import LogicalPath
 
+from . import _dtypes
 from ._private import MEDIA_TMP
 from .base_types.media import BatchableMedia, Media
 from .helper_types.bounding_boxes_2d import BoundingBoxes2D
@@ -687,3 +688,158 @@ class Image(BatchableMedia):
                 self._image = pil_image.open(self._path)
                 self._image.load()
         return self._image
+
+
+# Custom dtypes for typing system
+class _ImageFileType(_dtypes.Type):
+    name = "image-file"
+    legacy_names = ["wandb.Image"]
+    types = [Image]
+
+    def __init__(
+        self,
+        box_layers=None,
+        box_score_keys=None,
+        mask_layers=None,
+        class_map=None,
+        **kwargs,
+    ):
+        box_layers = box_layers or {}
+        box_score_keys = box_score_keys or []
+        mask_layers = mask_layers or {}
+        class_map = class_map or {}
+
+        if isinstance(box_layers, _dtypes.ConstType):
+            box_layers = box_layers._params["val"]
+        if not isinstance(box_layers, dict):
+            raise TypeError("box_layers must be a dict")
+        else:
+            box_layers = _dtypes.ConstType(
+                {layer_key: set(box_layers[layer_key]) for layer_key in box_layers}
+            )
+
+        if isinstance(mask_layers, _dtypes.ConstType):
+            mask_layers = mask_layers._params["val"]
+        if not isinstance(mask_layers, dict):
+            raise TypeError("mask_layers must be a dict")
+        else:
+            mask_layers = _dtypes.ConstType(
+                {layer_key: set(mask_layers[layer_key]) for layer_key in mask_layers}
+            )
+
+        if isinstance(box_score_keys, _dtypes.ConstType):
+            box_score_keys = box_score_keys._params["val"]
+        if not isinstance(box_score_keys, list) and not isinstance(box_score_keys, set):
+            raise TypeError("box_score_keys must be a list or a set")
+        else:
+            box_score_keys = _dtypes.ConstType(set(box_score_keys))
+
+        if isinstance(class_map, _dtypes.ConstType):
+            class_map = class_map._params["val"]
+        if not isinstance(class_map, dict):
+            raise TypeError("class_map must be a dict")
+        else:
+            class_map = _dtypes.ConstType(class_map)
+
+        self.params.update(
+            {
+                "box_layers": box_layers,
+                "box_score_keys": box_score_keys,
+                "mask_layers": mask_layers,
+                "class_map": class_map,
+            }
+        )
+
+    def assign_type(self, wb_type=None):
+        if isinstance(wb_type, _ImageFileType):
+            box_layers_self = self.params["box_layers"].params["val"] or {}
+            box_score_keys_self = self.params["box_score_keys"].params["val"] or []
+            mask_layers_self = self.params["mask_layers"].params["val"] or {}
+            class_map_self = self.params["class_map"].params["val"] or {}
+
+            box_layers_other = wb_type.params["box_layers"].params["val"] or {}
+            box_score_keys_other = wb_type.params["box_score_keys"].params["val"] or []
+            mask_layers_other = wb_type.params["mask_layers"].params["val"] or {}
+            class_map_other = wb_type.params["class_map"].params["val"] or {}
+
+            # Merge the class_ids from each set of box_layers
+            box_layers = {
+                str(key): set(
+                    list(box_layers_self.get(key, []))
+                    + list(box_layers_other.get(key, []))
+                )
+                for key in set(
+                    list(box_layers_self.keys()) + list(box_layers_other.keys())
+                )
+            }
+
+            # Merge the class_ids from each set of mask_layers
+            mask_layers = {
+                str(key): set(
+                    list(mask_layers_self.get(key, []))
+                    + list(mask_layers_other.get(key, []))
+                )
+                for key in set(
+                    list(mask_layers_self.keys()) + list(mask_layers_other.keys())
+                )
+            }
+
+            # Merge the box score keys
+            box_score_keys = set(list(box_score_keys_self) + list(box_score_keys_other))
+
+            # Merge the class_map
+            class_map = {
+                str(key): class_map_self.get(key, class_map_other.get(key, None))
+                for key in set(
+                    list(class_map_self.keys()) + list(class_map_other.keys())
+                )
+            }
+
+            return _ImageFileType(box_layers, box_score_keys, mask_layers, class_map)
+
+        return _dtypes.InvalidType()
+
+    @classmethod
+    def from_obj(cls, py_obj):
+        if not isinstance(py_obj, Image):
+            raise TypeError("py_obj must be a wandb.Image")
+        else:
+            if hasattr(py_obj, "_boxes") and py_obj._boxes:
+                box_layers = {
+                    str(key): set(py_obj._boxes[key]._class_labels.keys())
+                    for key in py_obj._boxes.keys()
+                }
+                box_score_keys = {
+                    key
+                    for val in py_obj._boxes.values()
+                    for box in val._val
+                    for key in box.get("scores", {}).keys()
+                }
+
+            else:
+                box_layers = {}
+                box_score_keys = set()
+
+            if hasattr(py_obj, "_masks") and py_obj._masks:
+                mask_layers = {
+                    str(key): set(
+                        py_obj._masks[key]._val["class_labels"].keys()
+                        if hasattr(py_obj._masks[key], "_val")
+                        else []
+                    )
+                    for key in py_obj._masks.keys()
+                }
+            else:
+                mask_layers = {}
+
+            if hasattr(py_obj, "_classes") and py_obj._classes:
+                class_set = {
+                    str(item["id"]): item["name"] for item in py_obj._classes._class_set
+                }
+            else:
+                class_set = {}
+
+            return cls(box_layers, box_score_keys, mask_layers, class_set)
+
+
+_dtypes.TypeRegistry.add(_ImageFileType)