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 BindingsEventWriter 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