burn/examples/pytorch-import
Guillaume Lagrange 4efc683df4
Upgrade to candle 0.4.1 (#1382)
* Fix python main entrypoint in book example

* Remove candle windows safeguards (#1178)

* Bump candle-core from 0.3.3 to 0.4.1

* Remove windows current known issue
2024-02-29 11:29:11 -06:00
..
model docs(book-&-examples): modify book and examples with new `prelude` module (#1372) 2024-02-28 13:25:25 -05:00
pytorch Add support for loading PyTorch `.pt` (weights/states) files directly to model's record (#1085) 2024-01-25 10:20:09 -05:00
src docs(book-&-examples): modify book and examples with new `prelude` module (#1372) 2024-02-28 13:25:25 -05:00
Cargo.toml [refactor] Move burn crates to their own crates directory (#1336) 2024-02-20 13:57:55 -05:00
README.md Add support for loading PyTorch `.pt` (weights/states) files directly to model's record (#1085) 2024-01-25 10:20:09 -05:00
build.rs Upgrade to candle 0.4.1 (#1382) 2024-02-29 11:29:11 -06:00

README.md

Import PyTorch Weights

This crate provides a simple example for importing PyTorch generated weights to Burn.

The .pt file is converted into a Burn consumable file (message pack format) using burn-import. The conversation is done in the build.rs file.

The model is separated into a sub-crate because build.rs needs for conversion and build.rs cannot import modules for the same crate.

Usage

cargo run -- 15

Output:

Finished dev [unoptimized + debuginfo] target(s) in 0.13s
    Running `burn/target/debug/onnx-inference 15`

Image index: 15
Success!
Predicted: 5
Actual: 5
See the image online, click the link below:
https://datasets-server.huggingface.co/assets/mnist/--/mnist/test/15/image/image.jpg