ferrotorch report

FERRO.md

Feedback on ferrotorch 0.6.2

Independent user feedback from porting a real PyTorch CNN (a 96×96 face-keypoint regression net: 3 conv blocks + BatchNorm + 3 FC, ~4.83 M params, masked-MSE loss, Adam + ReduceLROnPlateau, on-the-fly augmentation, DataLoader) from PyTorch to ferrotorch. This is written to help the project — both what's working and what got in the way. Everything below was hit while actually building and running the port, not from reading the README.

• ferrotorch version: 0.6.2
• toolchain: rustc 1.96.1 (Edition 2024), Linux x86_64, CPU-only (12 cores)
• date: June 2026
• workload: end-to-end train + eval + Python-weights parity check; full source at https://github.com/…/ferrotorch-examples/tree/main/face/ferrotorch (the face-ferrotorch crate).

Each item is tagged: [bug] correctness, [perf] performance, [api] ergonomics/consistency, [docs] documentation/examples, [pkg] packaging/deps. Verdicts are marked ✅ verified-by-running, 🔍 verified-by-reading-source, or ⚠️ inferred.

---

TL;DR

ferrotorch is a genuinely impressive effort: a PyTorch-shaped Rust DL framework where, in a single sitting, I could port a non-trivial CNN, run the training loop, and prove numerical parity with PyTorch to within float tolerance (Python-trained weights → ferrotorch model → 5.7 px test RMSE, matching the Python result). That it works at all, at this scope, in <4 months, is remarkable.

But three classes of issue will block adoption until fixed, in priority order: (1) CPU performance is impractically slow (~single-threaded), (2) the PyTorch .pt importer can't read real torch files, and (3) BatchNorm running stats silently fall out of state_dict(), producing wrong eval predictions after save/load. There are also a cluster of smaller API papercuts (no From<io::Error>, optimizer param-sync not public, mixed int types on shape ops, closure-only no_grad). Details and recommendations below.

---

What works well

• ✅ It builds clean and runs. Full default dep tree (438 crates) compiled in ~22 s on a fresh machine. The umbrella-re-export layout (ferrotorch::nn::*, ::optim::*, ::data::*, …) from the 28 sub-crates is the right call.
• ✅ The API really is PyTorch-shaped. Tensor<T: Float = f32>, Module<T> with forward(&self, &Tensor<T>) -> Result<Tensor<T>>, Sequential, Conv2d/Linear/BatchNorm2d/MaxPool2d/Dropout, Adam::new(params, cfg), DataLoader::new(Arc<D>, bs).shuffle(true).seed(42), nn::init::xavier_uniform, ReduceLROnPlateau. A PyTorch user is productive fast.
• ✅ Zero-panic, Result-everywhere API. Principled and consistently applied; ?-propagation through main() -> FerrotorchResult<()> was smooth.
• ✅ Numerical correctness on CPU. The parity test (load Python weights, forward in eval mode) reproduces the PyTorch result to a few px on 1410 test images. Autograd/Adam give a monotonically descending loss. This is the most important property and it holds.
• ✅ Thoughtful Module trait. Default state_dict/load_state_dict/ to_device/zero_grad/apply_to_parameters provided; as_any() downcast hook for layer-specific behavior; Parameter<T> derefs to Tensor<T>.
• ✅ StateDict<T> = HashMap<String, Tensor<T>> shared between nn and serialize — clean, no conversion at the serialization boundary.
• ✅ MetricScheduler vs LrScheduler split — ReduceLROnPlateau::step(&mut opt, metric_f64) is exactly right and distinct from metric-free schedulers.
• ✅ DataLoader::iter(epoch) seeding the per-epoch shuffle is a nice reproducibility primitive.
• ✅ Honest MNIST example that admits the cloned-params limitation and ships the sync helpers (more on this below — it should be a warning, not just a demo).

---

Issues

[bug] BatchNorm2d running stats are invisible to serialization — silent correctness loss 🔍✅

BatchNorm2d stores running_mean/running_var/num_batches_tracked as Mutex<Vec<f64>> (ferrotorch-nn/src/norm.rs) and overrides only as_any(), not buffers()/named_buffers()/state_dict()/load_state_dict(). Verified empirically: a BatchNorm2d's state_dict() returns exactly weight and bias — running_mean/running_var are absent.

Consequences (both reproduced):

1. save_pytorch(&model.state_dict(), "m.pt") writes a file that, on torch.load in Python, silently misses BN running stats → eval-mode predictions are wrong (running stats default to mean=0, var=1).
2. Module::load_state_dict(&sd, _) cannot restore them; the only path is m.as_any()?.downcast_ref::<BatchNorm2d<f32>>()?.set_running_mean(...).

For a layer as common as BatchNorm, this is a footgun that produces silent wrong answers at the serialization boundary — the worst failure mode.

Recommendation: expose BN running stats as Buffer<T> (preferred, so they flow through the default named_buffers()/state_dict()), or at minimum override named_buffers()/state_dict()/load_state_dict() on BatchNorm2d. Add a round-trip integration test that trains a BN model, saves, loads into a fresh model, and asserts identical eval predictions.

[bug] load_pytorch_state_dict cannot parse real torch .pt files ✅

serialize::load_pytorch_state_dict::<f32> round-trips its own save_pytorch output byte-identically (verified), but on an actual torch.save(state_dict, path) file produced by PyTorch 2.x it returns InvalidArgument { message: "no tensors found in pytorch state dict" }.

I disassembled the failing pickle: it is bog-standard protocol-2 — collections.OrderedDict → torch._utils _rebuild_tensor_v2 → torch.FloatStorage with BINPERSID referencing storage id "0" at location "cuda:0". The parser doesn't reconstruct tensors from this opcode sequence.

A related, compounding issue: find_pkl_name only tries ["archive/data.pkl", "data.pkl"] and find_data_prefix only ["archive/data/", "data/", ""] (ferrotorch-serialize/src/pytorch_import.rs). Modern torch.save writes the zip with a <filestem>/ prefix (my file used best_model/), so ferrotorch can't even locate the pickle. (I confirmed renaming to archive/ alone does not fix it — the deeper _rebuild_tensor_v2 parse still fails.)

Importing pretrained weights is the headline feature of a "PyTorch-shaped" framework; advertising load_pytorch_state_dict when it only reads your own writer erodes trust.

Recommendation:

• Scan the zip for any entry ending in /data.pkl (don't hardcode archive/).
• Teach the pickle parser the torch._utils._rebuild_tensor_v2 + persistent-id reconstruction path. Generate the test fixtures with a real torch.save in CI, not with ferrotorch's own writer.

(Workaround I used: convert .pt → .safetensors in Python with a stdlib-only writer; load_safetensors::<f32> reads it perfectly.)

[perf] CPU backend is effectively single-threaded — training is impractically slow ✅

On a 12-core box, training the 4.83 M-param CNN at batch 100 on 96×96 images used only ~124% CPU (~1.2 cores) and ran at ~160 s/epoch. For comparison, the same model trains an epoch in a few seconds in PyTorch on CPU. A 50-epoch run is ~2+ hours — unusable for iterative development.

The forward parity is correct, so this is purely a kernel-performance gap, likely concentrated in conv2d and the matmul behind Linear. target/release build; I confirmed release optimizations were on.

Recommendation: this is the #1 adoption blocker for anyone without a GPU. Prioritize, in order: (a) BLAS-backed matmul (matrixmultiply / a BLAS vendor), (b) rayon-parallel elementwise + conv kernels, (c) an im2col+GEMM conv path. Even 4–8× would move CPU from "unusable" to "tolerable for prototyping."

[api] Optimizer cloned-params sync is mandatory but not public API 🔍✅

The optimizer holds its own cloned Vec<Parameter<T>>. After loss.backward() writes grads onto the model's params, the optimizer's copies don't see them, and after opt.step() the model's params don't see the update. Users must copy two helpers from examples/train_mnist.rs:

fn sync_grads_to_optimizer(model: &dyn Module<T>, opt: &mut dyn Optimizer<T>) -> Result<()>
fn sync_params_from_optimizer(model: &mut dyn Module<T>, opt: &dyn Optimizer<T>) -> Result<()>

These are not pub. If a user forgets them, training silently does nothing (or trains on stale params) — another silent-failure footgun. The example's comment even says "A future API revision will use Arc-shared parameters."

Recommendation: until Arc-shared params land, (a) make these helpers pub API, and (b) put a loud # Warning on Optimizer::step's docstring describing the clone+sync contract. Better: have step read grads from the model params directly (pass &dyn Module to step, or have Adam::new take &mut-ish handles).

[api] No From<std::io::Error> (or From<csv::Error>, etc.) for FerrotorchError ✅

Every File::open(..)? and external-parser error needs .map_err(...). This is tedious for the extremely common "load a file" case. Compounded by InvalidArgument being a struct variant (InvalidArgument { message: String }, ferrotorch-core/src/error.rs:45) rather than a tuple variant, so the conversion is map_err(|e| FerrotorchError::InvalidArgument { message: e.to_string() }) every time.

Recommendation: add #[from] impls for at least std::io::Error and the common serialization crates you already depend on; consider a tuple-variant InvalidArgument(String) (or a generic Other(#[from] Box<dyn Error>)) for ergonomic wrapping.

[api] DataLoader has no general tensor collation ✅

DataLoader::iter(epoch) yields Vec<Sample> with no automatic stacking. default_collate handles bare Tensor<T>, default_collate_pair handles a 2-tuple, but there's nothing for N-tuples, and even the 2-tuple case only auto-stacks if you remember iter_collated + with_collate. For a (image, keypoints, mask) dataset I had to write a custom closure calling stack. PyTorch's default_collate handles arbitrary tuples/recursively.

Recommendation: a Collate trait (with a derive, or a blanket impl over tuples of Tensor) plugged in by default, so loader.iter(epoch) yields batched tensors out of the box.

[api] Mixed integer types on shape/dimension arguments ✅

reshape_t(&[isize]), flip_t(&[isize]), view(&[i64]), permute(&[usize]), transpose(usize, usize), narrow(usize, usize, usize). A user must remember which array type each op wants. PyTorch gets away with one (int).

Recommendation: standardize on one signed type (isize is the natural Rust choice since it allows -1) for all shape/dim args, with From impls to ease the &[usize]-vs-&[isize] boundary.

[api] pow_t takes f64, not T 🔍

Tensor::<f32>::pow_t(exponent: f64) — surprising that a method on Tensor<f32> doesn't take f32. Inconsistent with clamp_t(min: T, max: T). Minor, but it's the kind of thing that makes an API feel rough.

[api] no_grad is closure-only — no RAII guard ✅

no_grad(|| { ... }) forces borrowing everything used inside into a closure. When the model is also held &mut nearby, this fights the borrow checker (I hit this and fell back to the global set_grad_enabled(false) toggle).

Recommendation: also expose an RAII guard (let _g = no_grad_guard();) for parity with torch.no_grad() context-manager ergonomics. (And note in the docs that set_grad_enabled is process-global state — fine for single-threaded training, a hazard for multi-threaded data parallelism.)

[api] Constructor inconsistency between layers 🔍

Conv2d::<T>::new(..) -> Result<Self>, Linear::<T>::new(..) -> Result<Self>, Dropout::<T>::new(p) -> Result<Self> (all fallible, generic), but MaxPool2d::new(..) -> Self, LeakyReLU::new(..) -> Self, Flatten::new(..) -> Self (infallible, non-generic). The mix is mildly jarring when constructing a Sequential (some layers need ?, some don't; some need ::<f32>, some don't). Non-generic stateless layers are fine, but consider making the fallible ones infallible where the inputs can't actually fail (e.g. Dropout::new(0.2) only fails on out-of-range p, which could be a debug_assert instead).

[api] gather/index_select ergonomics 🔍

The free-fn gather takes &[usize] + a shape; the method forms are generic over IntElement and want an IntTensor<I>. For the common case "reorder along an axis with a Vec<usize> of indices" (e.g. a keypoint flip permutation), there's no convenient path; I did host-side indexing instead. Worth a Tensor::index_axis_simple(dim, &[usize]) or similar.

[docs] No end-to-end CNN example ✅

The only training example is a 3-layer MLP (examples/train_mnist.rs). There's no example wiring Conv2d + BatchNorm2d + MaxPool2d + Flatten + DataLoader + optimizer + backward. I had to discover the layer constructor signatures and the BN-buffer behavior by reading source. A examples/train_cnn.rs (even CIFAR-ish) would massively lower the barrier and would have caught several of the issues above internally.

[docs] #[derive(Module)] contract is underspecified 🔍

The derive requires a field literally named training: bool plus #[param] / #[submodule] / #[skip] attributes, and — critically — does not generate forward for custom topologies, so any model with reused layers or non-linear wiring ends up hand-implementing Module anyway. That should be up front in the docs, with a worked hand-impl example (which is the more common real-world case).

[pkg] Default feature set is heavy ✅

A plain ferrotorch = "0.6.2" (no features) pulls 438 crates including tokenizers, minijinja, ureq, cranelift-jit (the JIT!), tracing, safetensors, etc. For someone who just wants to train a small CNN on CPU, this is a lot of compile time and binary size for unused functionality.

Recommendation: move jit, hub, tokenize, distributions, profiler, ml (and friends) behind opt-in features; keep the default to the core+nn+optim+data+serialize+train needed for a basic train loop. This also shrinks the attack surface and the dep-update churn.

[bug?] rand::Rng::gen is a reserved keyword in Edition 2024 — update examples/docs ⚠️

Minor, but anything showing rng.gen::<f32>() breaks under Edition 2024 (gen is reserved for gen-blocks). ferrotorch itself requires Edition 2024, so any user-facing sample code using .gen() will frustrate newcomers. Show gen_range(..) or r#gen(..) instead. (This isn't ferrotorch's bug, but it's ferrotorch's docs to keep current.)

---

Maturity / testing observations

• The README claims 9800+ tests; the test suite is visibly large (many conformance_*, audit_*, divergence_*, _probe_* files). Yet I found real bugs in core paths — pickle import of real torch files, BN serialization, CPU performance — that this volume of tests didn't catch. This suggests the tests are heavily structural/inventory-style (asserting that symbols exist and that internal round-trips work) rather than integration tests against real external artifacts. The pickle parser passing its own writer's output but failing torch's is the canonical example.
• Recommendation: add a small set of "real-world" integration tests: (a) check in a .pt generated by a reference torch.save and assert it loads; (b) train a BN CNN for a few steps, save, load into a fresh model, assert byte/elem-equal predictions; (c) a CNN training smoke that runs >1 epoch.
• The development style (REQ tables in doc-comments, Closes #N, _probe_* / divergence_*_reattempt.rs naming) reads as heavily AI-assisted/audit-driven. That's a legitimate workflow and it produced a lot of working code, but it trends toward APIs that are self-consistent but externally unvalidated. Pairing it with real-world integration tests would close the gap and is the single highest-leverage quality investment.
• Adoption signals are early (0 stars, ~500 downloads at time of writing), so first impressions matter enormously — the three top issues above (CPU perf, real .pt import, BN serialization) are exactly the things a first-time user on a non-trivial model will hit.

---

What I did not exercise

In the interest of honesty: I did not test the GPU/mps/xpu/cubecl backends, the JIT/tracing, distributed training, RNN/Transformer/attention modules, ONNX export, GGUF, the Learner high-level training API, or vision models beyond Mnist::synthetic. My feedback covers core tensor ops, nn (conv/bn/pool/ linear/dropout/activations), optim (Adam + ReduceLROnPlateau), data (Dataset + DataLoader), and serialize (save_pytorch, load_pytorch_state_dict, load_safetensors) on CPU.

---

Prioritized recommendations

1. CPU performance — parallelize kernels + BLAS matmul/im2col. (#1 adoption blocker.)
2. BatchNorm running stats in state_dict() — expose as Buffer<T>; add a train→save→load→eval parity test.
3. Real .pt import — handle _rebuild_tensor_v2 + persistent-id and arbitrary zip prefixes; CI-generate fixtures with torch.save.
4. Public optimizer param-sync helpers + a # Warning on Optimizer::step.
5. Slim default features (move JIT/hub/tokenize/distributions behind opt-in).
6. From<io::Error> + tuple-variant InvalidArgument for ergonomic file IO.
7. A CNN training example (examples/train_cnn.rs) — would have caught 1–3 internally.
8. Auto-collation in DataLoader (tuple-aware Collate trait).
9. Standardize shape/dim int types (one signed type).
10. RAII no_grad guard alongside the closure form.

Thanks for the work — the bones here are good, and the parity result proves the numerics are sound. Fixing the top three would make ferrotorch genuinely usable for real work, not just demos.

so in summary: ferrotorch rocks ❤️ 🔥