Skip to content

PRATHAM777P/tensorflow-rs

Folders and files

NameName
Last commit message
Last commit date

Latest commit

ย 

History

33 Commits
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 
ย 

Repository files navigation

tensorflow-rs

A blazing-fast, pure-Rust library for reading and writing TensorBoard event files โ€” with first-class Python bindings via PyO3.

CI Crates.io License: Apache-2.0 Rust Edition PRs Welcome

Getting Started ยท API Reference ยท Examples ยท Contributing


๐Ÿ” What is tensorflow-rs?

tensorflow-rs lets you write and read TensorBoard .tfevents.* files directly from Rust โ€” no TensorFlow, no Python required (unless you want the bindings!).

TensorBoard is the standard visualization toolkit in deep learning โ€” displaying training metrics like loss, accuracy, and gradients over time. This library generates those log files natively from Rust (or Python via PyO3), making it perfect for:

Use Case Description
๐Ÿฆ€ Rust ML Training Works with Candle, tch-rs, and any Rust ML framework
๐Ÿ Python Pipelines Rust-speed logging with a clean Python API via PyO3/maturin
๐Ÿ“Š Experiment Tracking Emit TensorBoard-compatible logs from any custom tool

โœจ Features

  • ๐Ÿ“ˆ Scalars โ€” log training loss, accuracy, learning rate, and any float metric
  • ๐Ÿ“Š Histograms โ€” visualize weight and gradient distributions over training
  • ๐Ÿ–ผ๏ธ Images โ€” log generated or debug images at any training step
  • ๐Ÿ”Š Audio โ€” log WAV audio samples directly from raw PCM data
  • ๐Ÿงฎ Tensors โ€” log typed tensors (f32, f64, i32, i64)
  • ๐Ÿ“– Event Reader โ€” parse existing .tfevents.* files back to structured Rust types
  • ๐Ÿ Python Bindings โ€” EventWriter and EventReader exposed as Python classes
  • โœ… CRC-Verified I/O โ€” every TFRecord is checksum-validated for correctness
  • ๐Ÿšซ Zero TF Dependency โ€” pure Rust with protobuf via prost
  • ๐Ÿ–ฅ๏ธ Cross-Platform โ€” tested on Linux, macOS, and Windows

๐Ÿ“ Project Structure

tensorflow-rs/
โ”œโ”€โ”€ Cargo.toml                        # Workspace root
โ”‚
โ”œโ”€โ”€ tboard/                           # Core Rust library crate
โ”‚   โ”œโ”€โ”€ Cargo.toml
โ”‚   โ”œโ”€โ”€ build.rs                      # Compiles .proto files at build time
โ”‚   โ””โ”€โ”€ src/
โ”‚       โ”œโ”€โ”€ lib.rs                    # Public API + CRC helper
โ”‚       โ”œโ”€โ”€ writer.rs                 # EventWriter: creates & writes tfevents files
โ”‚       โ”œโ”€โ”€ reader.rs                 # SummaryReader: parses tfevents as an iterator
โ”‚       โ”œโ”€โ”€ error.rs                  # Custom error types + bail! macro
โ”‚       โ”œโ”€โ”€ wave.rs                   # PCM โ†’ WAV encoder (audio logging)
โ”‚       โ”œโ”€โ”€ event.proto               # TF Event protobuf schema
โ”‚       โ”œโ”€โ”€ summary.proto             # TF Summary protobuf schema
โ”‚       โ”œโ”€โ”€ tensor.proto              # TF Tensor protobuf schema
โ”‚       โ”œโ”€โ”€ histogram.proto           # TF Histogram protobuf schema
โ”‚       โ””โ”€โ”€ *.proto                   # Other TF-compatible proto definitions
โ”‚   โ””โ”€โ”€ examples/
โ”‚       โ””โ”€โ”€ read.rs                   # CLI: read & pretty-print a tfevents file
โ”‚
โ”œโ”€โ”€ tboard-pyo3/                      # Python bindings (PyO3 + maturin)
โ”‚   โ”œโ”€โ”€ Cargo.toml
โ”‚   โ”œโ”€โ”€ pyproject.toml                # Python packaging config
โ”‚   โ””โ”€โ”€ src/
โ”‚       โ””โ”€โ”€ lib.rs                    # Python-exposed EventWriter + EventReader
โ”‚
โ”œโ”€โ”€ basics.py                         # Quick Python demo script
โ”œโ”€โ”€ rustfmt.toml                      # Rust formatting config
โ””โ”€โ”€ .github/workflows/
    โ””โ”€โ”€ rust-ci.yml                   # CI: check, test, fmt, clippy โ€” 3 OSes ร— 2 toolchains

๐Ÿš€ Getting Started

Prerequisites

Tool Purpose
Rust (stable or nightly) Core build toolchain
protoc (install guide) Protocol Buffer compiler
Python โ‰ฅ 3.8 + maturin Python bindings only

Rust Usage

Add tboard to your Cargo.toml:

[dependencies]
tboard = "0.1.1"

Write training metrics:

use tboard::EventWriter;

fn main() -> tboard::Result<()> {
    let mut writer = EventWriter::create("./logs")?;

    for step in 0..1000i64 {
        let loss = 1.0 / (1.0 + step as f32 * 0.01);
        writer.write_scalar(step, "loss/train", loss)?;
    }

    writer.flush()?;
    Ok(())
}

Then visualize in TensorBoard:

tensorboard --logdir ./logs

Python Usage

Build and install the Python wheel with maturin:

cd tboard-pyo3
pip install maturin
maturin develop

Then use it in Python:

import tboard, math

# --- Writing ---
tb = tboard.EventWriter("/tmp/my-experiment")
for step in range(10000):
    tb.add_scalar("loss",     math.exp(-step * 0.001), step)
    tb.add_scalar("accuracy", 1 - math.exp(-step * 0.001), step)
tb.flush()
print("Wrote to:", tb.filename)

# --- Reading ---
reader = tboard.EventReader(tb.filename)
for event in reader:
    print(event)

๐Ÿ“š API Reference

EventWriter

The main struct for writing TensorBoard event files.

Rust API
// Create a new writer in a log directory
let mut writer = EventWriter::create("./logdir")?;

// Write a scalar
writer.write_scalar(step: i64, tag: &str, value: f32)?;

// Write a histogram
writer.write_histo(step, tag, min, max, num, sum, sum_squares, bucket, bucket_limit)?;

// Write an image (encoded bytes, e.g. PNG)
writer.write_image(step, tag, width, height, colorspace, encoded_image_string)?;

// Write audio from raw PCM samples (auto-encodes to WAV)
writer.write_pcm_as_wav(step, tag, pcm_data: &[f32], sample_rate: u32)?;

// Write a typed tensor
writer.write_tensor::<f32>(step, tag, values: Vec<f32>)?;

// Flush buffered data
writer.flush()?;
Python API
writer = tboard.EventWriter(logdir: str, on_error: str = "raise")
# on_error: "raise" (default) | "log" (print errors, don't crash)

writer.add_scalar(tag: str, scalar_value: float, global_step: int = 0)
writer.add_audio(tag: str, pcm_data: List[float], sample_rate: int, global_step: int = 0)
writer.flush()

writer.logdir    # the log directory
writer.filename  # full path of the current event file

SummaryReader

An iterator-based reader for existing TensorBoard event files.

Rust API
use tboard::SummaryReader;
use std::fs::File;

let reader = SummaryReader::new(File::open("events.out.tfevents....")?);

for event in reader {
    let event = event?;
    println!("step={} wall_time={:.3}", event.step, event.wall_time);
}
Python API
reader = tboard.EventReader(filename: str)

for event in reader:
    # event is a dict with keys:
    # - "step":      int
    # - "wall_time": float
    # - "kind":      str   ("summary", "file_version", "graph_def", ...)
    # - "what":      list  (for summaries: [{"tag": str, "value": float}])
    print(event)

๐Ÿ’ก Examples

Rust CLI Reader

Read and pretty-print any .tfevents file:

cargo run --example read -- path/to/events.out.tfevents.XXXXXXXXXX

Sample output:

2024-02-04 12:00:01 UTC      step:        0     loss/train: 1.0000
2024-02-04 12:00:01 UTC      step:        1     loss/train: 0.9901
2024-02-04 12:00:01 UTC      step:        2     loss/train: 0.9803

Python Quick Demo

python basics.py

Writes 100,000 scalar steps of sin(step ร— 1e-4) to /tmp/test-event-writer, then reads them back and prints each event.


๐Ÿ”จ Building from Source

# Clone the repo
git clone https://github.com/PRATHAM777P/tensorflow-rs.git
cd tensorflow-rs

# Build (protoc must be installed)
cargo build --release

# Run tests
cargo test

# Run Clippy lints
cargo clippy -- -D warnings

# Check formatting
cargo fmt --all -- --check

Python bindings:

cd tboard-pyo3
maturin develop          # install into current virtualenv
maturin build --release  # build a distributable .whl wheel

๐Ÿ–ฅ๏ธ CI / Supported Platforms

CI runs on every push and pull request across 6 matrix targets:

OS Stable Nightly
Ubuntu latest โœ… โœ…
macOS latest โœ… โœ…
Windows latest โœ… โœ…

Jobs: check ยท test ยท rustfmt ยท clippy


๐Ÿ“œ License

Licensed under the Apache License, Version 2.0.


๐Ÿ‘ค Author

Prathamesh Penshanwar โ€” @PRATHAM777P

Built with โค๏ธ in Rust. Contributions, issues, and PRs are very welcome!


Proto definitions sourced from the TensorBoard protobuf directory โ€” used under the Apache 2.0 License.

About

A blazing-fast pure-Rust library for reading and writing TensorBoard event files (.tfevents), with Python bindings via PyO3 no TensorFlow dependency required.

Topics

Resources

License

Security policy

Stars

Watchers

Forks

Contributors