wookayin / expt

Experiment. Plot. Tabulate.
MIT License
68 stars 6 forks source link

Add a native, fast tensorboard eventfile reader based on rustboard #13

Closed wookayin closed 1 year ago

wookayin commented 1 year ago

This commit adds a rust extension module exposed as expt._internal, implementing a class TensorboardEventFileReader` which can read tensorboard eventfiles and logdir at a blazingly fast speed.

Rust implementation:

The core logic is borrowed from tensorboard's rust backend, a.k.a. "rustboard". For now, we extract all the scalar summary entires in the form of DataFrame-like dictionary, and convert it to a pandas DataFrame.

Build system:

Build of rust native extension (expt._internal) is based on setuptools-rust. We also considered maturin for the build tool, but since expt is still a python-major project we can try a simpler build tool such as setuptools-rust for the time being.

Performance:

Parsing tensorboard eventfiles is roughly 20x faster than the previous python implementation (which uses tensorflow's protobuf-based native EventFile loader).

See fixtures/README.md for preliminary benchmark results.