diff --git a/clu/metric_writers/async_writer.py b/clu/metric_writers/async_writer.py index 3fa28a5..c5b1dbd 100644 --- a/clu/metric_writers/async_writer.py +++ b/clu/metric_writers/async_writer.py @@ -23,7 +23,7 @@ from collections.abc import Mapping, Sequence import contextlib -from typing import Any, Optional, Union +from typing import Any, Optional from clu import asynclib @@ -31,7 +31,6 @@ from clu.metric_writers import multi_writer import wrapt - Array = interface.Array Scalar = interface.Scalar @@ -97,44 +96,21 @@ def write_videos(self, step: int, videos: Mapping[str, Array]): @_wrap_exceptions def write_audios( - self, step: int, audios: Mapping[str, Array], *, sample_rate: int - ): + self, step: int, audios: Mapping[str, Array], *, sample_rate: int): self._pool(self._writer.write_audios)( - step=step, audios=audios, sample_rate=sample_rate - ) + step=step, audios=audios, sample_rate=sample_rate) @_wrap_exceptions def write_texts(self, step: int, texts: Mapping[str, str]): self._pool(self._writer.write_texts)(step=step, texts=texts) @_wrap_exceptions - def write_histograms( - self, - step: int, - arrays: Mapping[str, Array], - num_buckets: Optional[Mapping[str, int]] = None, - ): + def write_histograms(self, + step: int, + arrays: Mapping[str, Array], + num_buckets: Optional[Mapping[str, int]] = None): self._pool(self._writer.write_histograms)( - step=step, arrays=arrays, num_buckets=num_buckets - ) - - @_wrap_exceptions - def write_pointcloud( - self, - step: int, - point_clouds: Mapping[str, Array], - *, - point_colors: Optional[Array] = None, - configs: Optional[ - Mapping[str, Union[str, int, float, bool, None]] - ] = None, - ): - self._pool(self._writer.write_pointcloud)( - step=step, - point_clouds=point_clouds, - point_colors=point_colors, - configs=configs, - ) + step=step, arrays=arrays, num_buckets=num_buckets) @_wrap_exceptions def write_hparams(self, hparams: Mapping[str, Any]): diff --git a/clu/metric_writers/async_writer_test.py b/clu/metric_writers/async_writer_test.py index c3d10a5..736fa6c 100644 --- a/clu/metric_writers/async_writer_test.py +++ b/clu/metric_writers/async_writer_test.py @@ -72,27 +72,6 @@ def test_write_videos(self): self.sync_writer.write_videos.assert_called_with(4, {"input_videos": mock.ANY}) - def test_write_pointcloud(self): - point_clouds = np.random.normal(0, 1, (1, 1024, 3)).astype(np.float32) - point_colors = np.random.uniform(0, 1, (1, 1024, 3)).astype(np.float32) - config = { - "material": "PointCloudMaterial", - "size": 0.09, - } - self.writer.write_pointcloud( - step=0, - point_clouds={"pcd": point_clouds}, - point_colors={"pcd": point_colors}, - configs={"config": config}, - ) - self.writer.flush() - self.sync_writer.write_pointcloud.assert_called_with( - step=0, - point_clouds={"pcd": mock.ANY}, - point_colors={"pcd": mock.ANY}, - configs={"config": mock.ANY}, - ) - def test_write_texts(self): self.writer.write_texts(4, {"samples": "bla"}) self.writer.flush() diff --git a/clu/metric_writers/interface.py b/clu/metric_writers/interface.py index 1c73846..c497bfe 100644 --- a/clu/metric_writers/interface.py +++ b/clu/metric_writers/interface.py @@ -153,28 +153,6 @@ def write_histograms(self, of the MetricWriter. """ - @abc.abstractmethod - def write_pointcloud( - self, - step: int, - point_clouds: Mapping[str, Array], - *, - point_colors: Optional[Mapping[str, Array]] = None, - configs: Optional[ - Mapping[str, Union[str, int, float, bool, None]] - ] = None, - ): - """Writes point cloud summaries. - - Args: - step: Step at which the point cloud was generated. - point_clouds: Mapping from point clouds key to point cloud of shape [N, 3] - array of point coordinates. - point_colors: Mapping from point colors key to [N, 3] array of point - colors. - configs: A dictionary of configuration options for the point cloud. - """ - @abc.abstractmethod def write_hparams(self, hparams: Mapping[str, Any]): """Write hyper parameters. diff --git a/clu/metric_writers/logging_writer.py b/clu/metric_writers/logging_writer.py index 2e97cdd..6e8a28d 100644 --- a/clu/metric_writers/logging_writer.py +++ b/clu/metric_writers/logging_writer.py @@ -15,7 +15,7 @@ """MetricWriter that writes all values to INFO log.""" from collections.abc import Mapping -from typing import Any, Optional, Union +from typing import Any, Optional from absl import logging from clu.metric_writers import interface @@ -77,29 +77,6 @@ def write_histograms(self, self._collection_str, key, _get_histogram_as_string(histo, bins)) - def write_pointcloud( - self, - step: int, - point_clouds: Mapping[str, Array], - *, - point_colors: Optional[Mapping[str, Any]] = None, - configs: Optional[ - Mapping[str, Union[str, int, float, bool, None]] - ] = None, - ): - logging.info( - "[%d]%s Got point clouds: %s, point_colors: %s, configs: %s.", - step, - self._collection_str, - {k: v.shape for k, v in point_clouds.items()}, - ( - {k: v.shape for k, v in point_colors.items()} - if point_colors is not None - else None - ), - configs, - ) - def write_hparams(self, hparams: Mapping[str, Any]): logging.info("[Hyperparameters]%s %s", self._collection_str, hparams) diff --git a/clu/metric_writers/logging_writer_test.py b/clu/metric_writers/logging_writer_test.py index 9cc0b65..9683338 100644 --- a/clu/metric_writers/logging_writer_test.py +++ b/clu/metric_writers/logging_writer_test.py @@ -80,42 +80,14 @@ def test_write_histogram(self): "INFO:absl:[4] Histogram for 'c' = {[-0.4, 0.6]: 5}", ]) - def test_write_pointcloud(self): - point_clouds = np.random.normal(0, 1, (1, 1024, 3)).astype(np.float32) - point_colors = np.random.uniform(0, 1, (1, 1024, 3)).astype(np.float32) - config = { - "material": "PointCloudMaterial", - "size": 0.09, - } - with self.assertLogs(level="INFO") as logs: - self.writer.write_pointcloud( - step=4, - point_clouds={"pcd": point_clouds}, - point_colors={"pcd": point_colors}, - configs={"configs": config}, - ) - self.assertEqual( - logs.output, - [ - "INFO:absl:[4] Got point clouds: {'pcd': (1, 1024, 3)}," - " point_colors: {'pcd': (1, 1024, 3)}, configs: {'configs':" - " {'material': 'PointCloudMaterial', 'size': 0.09}}." - ], - ) - def test_write_hparams(self): with self.assertLogs(level="INFO") as logs: self.writer.write_hparams({"learning_rate": 0.1, "batch_size": 128}) - self.assertEqual( - logs.output, - [ - "INFO:absl:[Hyperparameters] {'learning_rate': 0.1, 'batch_size':" - " 128}" - ], - ) + self.assertEqual(logs.output, [ + "INFO:absl:[Hyperparameters] {'learning_rate': 0.1, 'batch_size': 128}" + ]) def test_collection(self): - writer = logging_writer.LoggingWriter(collection="train") writer = logging_writer.LoggingWriter(collection="train") with self.assertLogs(level="INFO") as logs: writer.write_scalars(0, {"a": 3, "b": 0.15}) diff --git a/clu/metric_writers/multi_writer.py b/clu/metric_writers/multi_writer.py index bf6a9b9..fc300e2 100644 --- a/clu/metric_writers/multi_writer.py +++ b/clu/metric_writers/multi_writer.py @@ -15,7 +15,7 @@ """MetricWriter that writes to multiple MetricWriters.""" from collections.abc import Mapping, Sequence -from typing import Any, Optional, Union +from typing import Any, Optional from clu.metric_writers import interface @@ -57,27 +57,13 @@ def write_texts(self, step: int, texts: Mapping[str, str]): for w in self._writers: w.write_texts(step, texts) - def write_histograms( - self, - step: int, - arrays: Mapping[str, Array], - num_buckets: Optional[Mapping[str, int]] = None): + def write_histograms(self, + step: int, + arrays: Mapping[str, Array], + num_buckets: Optional[Mapping[str, int]] = None): for w in self._writers: w.write_histograms(step, arrays, num_buckets) - def write_pointcloud( - self, - step: int, - point_clouds: Mapping[str, Array], - *, - point_colors: Optional[Mapping[str, Array]] = None, - configs: Optional[Mapping[str, Union[str, int, float, bool, None]]] = None - ): - for w in self._writers: - w.write_pointcloud( - step, point_clouds, point_colors=point_colors, configs=configs - ) - def write_hparams(self, hparams: Mapping[str, Any]): for w in self._writers: w.write_hparams(hparams) diff --git a/clu/metric_writers/multi_writer_test.py b/clu/metric_writers/multi_writer_test.py index 350ea9a..5f53b44 100644 --- a/clu/metric_writers/multi_writer_test.py +++ b/clu/metric_writers/multi_writer_test.py @@ -18,7 +18,6 @@ from clu.metric_writers import interface from clu.metric_writers import multi_writer -import numpy as np import tensorflow as tf @@ -49,29 +48,6 @@ def test_write_scalars(self): ]) w.flush.assert_called() - def test_write_pointcloud(self): - point_clouds = np.random.normal(0, 1, (1, 1024, 3)).astype(np.float32) - point_colors = np.random.uniform(0, 1, (1, 1024, 3)).astype(np.float32) - config = { - "material": "PointCloudMaterial", - "size": 0.09, - } - self.writer.write_pointcloud( - step=0, - point_clouds={"pcd": point_clouds}, - point_colors={"pcd": point_colors}, - configs={"config": config}, - ) - self.writer.flush() - for w in self.writers: - w.write_pointcloud.assert_called_with( - step=0, - point_clouds={"pcd": point_clouds}, - point_colors={"pcd": point_colors}, - configs={"config": config}, - ) - w.flush.assert_called() - if __name__ == "__main__": tf.test.main() diff --git a/clu/metric_writers/tf/summary_writer.py b/clu/metric_writers/tf/summary_writer.py index 4833e69..081e8e9 100644 --- a/clu/metric_writers/tf/summary_writer.py +++ b/clu/metric_writers/tf/summary_writer.py @@ -19,7 +19,7 @@ """ from collections.abc import Mapping -from typing import Any, Optional, Union +from typing import Any, Optional from absl import logging @@ -31,7 +31,6 @@ with epy.lazy_imports(): # pylint: disable=g-import-not-at-top from tensorboard.plugins.hparams import api as hparams_api - from tensorboard.plugins.mesh import summary as mesh_summary # pylint: disable=line-too-long # pylint: enable=g-import-not-at-top @@ -98,28 +97,6 @@ def write_histograms( buckets = None if num_buckets is None else num_buckets.get(key) tf.summary.histogram(key, value, step=step, buckets=buckets) - def write_pointcloud( - self, - step: int, - point_clouds: Mapping[str, Array], - *, - point_colors: Optional[Mapping[str, Array]] = None, - configs: Optional[ - Mapping[str, Union[str, int, float, bool, None]] - ] = None, - ): - with self._summary_writer.as_default(): - for key, vertices in point_clouds.items(): - colors = None if point_colors is None else point_colors.get(key) - config = None if configs is None else configs.get(key) - mesh_summary.mesh( - key, - vertices=vertices, - colors=colors, - step=step, - config_dict=config, - ) - def write_hparams(self, hparams: Mapping[str, Any]): with self._summary_writer.as_default(): hparams_api.hparams(dict(utils.flatten_dict(hparams))) diff --git a/clu/metric_writers/tf/summary_writer_test.py b/clu/metric_writers/tf/summary_writer_test.py index ce156f4..bb8eb87 100644 --- a/clu/metric_writers/tf/summary_writer_test.py +++ b/clu/metric_writers/tf/summary_writer_test.py @@ -83,21 +83,6 @@ def _load_hparams(logdir: str): return hparams -def _load_pointcloud_data(logdir: str): - """Loads pointcloud summaries from events in a logdir.""" - paths = tf.io.gfile.glob(os.path.join(logdir, "events.out.tfevents.*")) - data = collections.defaultdict(dict) - for path in paths: - for event in tf.compat.v1.train.summary_iterator(path): - for value in event.summary.value: - if value.metadata.plugin_data.plugin_name == "mesh": - if "config" not in value.tag: - data[event.step][value.tag] = tf.make_ndarray(value.tensor) - else: - data[event.step][value.tag] = value.metadata.plugin_data.content - return data - - class SummaryWriterTest(tf.test.TestCase): def setUp(self): @@ -157,24 +142,6 @@ def test_write_histograms(self): ] self.assertAllClose(data["b"], ([0, 2], expected_histograms_b)) - def test_write_pointcloud(self): - point_clouds = np.random.normal(0, 1, (1, 1024, 3)).astype(np.float32) - point_colors = np.random.uniform(0, 1, (1, 1024, 3)).astype(np.float32) - config = { - "material": "PointCloudMaterial", - "size": 0.09, - } - self.writer.write_pointcloud( - step=0, - point_clouds={"pcd": point_clouds}, - point_colors={"pcd": point_colors}, - configs={"config": config}, - ) - self.writer.flush() - data = _load_pointcloud_data(self.logdir) - self.assertAllClose(data[0]["pcd_VERTEX"], point_clouds) - self.assertAllClose(data[0]["pcd_COLOR"], point_colors) - def test_hparams(self): self.writer.write_hparams(dict(batch_size=512, num_epochs=90)) hparams = _load_hparams(self.logdir) diff --git a/clu/metric_writers/torch_tensorboard_writer.py b/clu/metric_writers/torch_tensorboard_writer.py index 279dce1..77ce35c 100644 --- a/clu/metric_writers/torch_tensorboard_writer.py +++ b/clu/metric_writers/torch_tensorboard_writer.py @@ -19,7 +19,7 @@ """ from collections.abc import Mapping -from typing import Any, Optional, Union +from typing import Any, Optional from absl import logging @@ -79,22 +79,6 @@ def write_histograms(self, self._writer.add_histogram( tag, values, global_step=step, bins="auto", max_bins=bins) - def write_pointcloud( - self, - step: int, - point_clouds: Mapping[str, Array], - *, - point_colors: Optional[Mapping[str, Array]] = None, - configs: Optional[ - Mapping[str, Union[str, int, float, bool, None]] - ] = None, - ): - logging.log_first_n( - logging.WARNING, - "TorchTensorBoardWriter does not support writing point clouds.", - 1, - ) - def write_hparams(self, hparams: Mapping[str, Any]): self._writer.add_hparams(hparams, {})