wehub/tracel-ai--burn
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:45:29 +08:00
commit 5b45ce7982
2056 changed files with 390094 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.cargo/audit.toml
+32
View File
@@ -0,0 +1,32 @@
# Audit config file
#
# It may be located in the user home (`~/.cargo/audit.toml`) or in the project
# root (`.cargo/audit.toml`).
#
# All of the options which can be passed via CLI arguments can also be
# permanently specified in this file.
[advisories]
ignore = [
"RUSTSEC-2024-0436", # Paste used to generate macro, should be removed at some point.
"RUSTSEC-2025-0119", # `number_prefix` used by `tokenizers`, only in the examples.
"RUSTSEC-2025-0141", # `bincode` is no longer maintained.
"RUSTSEC-2024-0388", # `derivative` dependancy in the DQN example is unmaintained.
"RUSTSEC-2026-0194", # `quick-xml` <0.41.0, pinned by polars → object_store 0.13.2; no upgrade path yet.
"RUSTSEC-2026-0195", # `quick-xml` <0.41.0, pinned by polars → object_store 0.13.2; no upgrade path yet.
] # advisory IDs to ignore e.g. ["RUSTSEC-2019-0001", ...]
informational_warnings = [
"unmaintained",
] # warn for categories of informational advisories
severity_threshold = "low" # CVSS severity ("none", "low", "medium", "high", "critical")
# Output Configuration
[output]
deny = ["unmaintained"] # exit on error if unmaintained dependencies are found
format = "terminal" # "terminal" (human readable report) or "json"
quiet = false # Only print information on error
show_tree = true # Show inverse dependency trees along with advisories (default: true)
[yanked]
enabled = true # Warn for yanked crates in Cargo.lock (default: true)
update_index = true # Auto-update the crates.io index (default: true)
.cargo/config.toml
+3
View File
@@ -0,0 +1,3 @@
[alias]
xtask = "run --target-dir target/xtask --color always --package xtask --bin xtask --"
run-checks = "xtask -c all validate --release"
.github/ISSUE_TEMPLATE/bug_report.md
+40
View File
@@ -0,0 +1,40 @@
---
name: Bug report
about: Create a report to help us improve
title: ''
labels: ''
assignees: ''
---
**Describe the bug**
<!-- A clear and concise description of what the bug is. -->
**To Reproduce**
<!--
Steps to reproduce the behavior:
1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error
-->
**Expected behavior**
<!-- A clear and concise description of what you expected to happen. -->
**Screenshots**
<!-- If applicable, add screenshots to help explain your problem. -->
**Desktop (please complete the following information):**
- OS: [e.g. iOS]
- Browser [e.g. chrome, safari]
- Version [e.g. 22]
**Smartphone (please complete the following information):**
- Device: [e.g. iPhone6]
- OS: [e.g. iOS8.1]
- Browser [e.g. stock browser, safari]
- Version [e.g. 22]
**Additional context**
<!-- Add any other context about the problem here. -->
.github/ISSUE_TEMPLATE/doc_request.md
+10
View File
@@ -0,0 +1,10 @@
---
name: Documentation request
about: Flag incoherent or missing documentation, including use case examples.
title: ''
labels: ''
assignees: ''
---
<!-- Please search existing issues to avoid creating duplicates -->
.github/ISSUE_TEMPLATE/feature_request.md
+28
View File
@@ -0,0 +1,28 @@
---
name: Feature request
about: Suggest an idea for this project
title: ''
labels: ''
assignees: ''
---
<!-- Please search existing issues to avoid creating duplicates -->
### Feature description
<!-- Describe the feature you'd like -->
### Feature motivation
<!-- Why do you want this? -->
### (Optional) Suggest a Solution
<!--
How do you think we should implement this feature?
Things to address include:
* Details of the technical implementation
* Tradeoffs made in design decisions
* Caveats and considerations for the future
-->
.github/PULL_REQUEST_TEMPLATE/template.md
+12
View File
@@ -0,0 +1,12 @@
* **Please check if the PR fulfills these requirements**
- [ ] The commit message follows our guidelines
- [ ] Docs have been added / updated (for bug fixes / features)
* **What kind of change does this PR introduce?** (Bug fix, feature, docs update, ...)
* **Does this PR introduce a breaking change?** (What changes might users need to make in their application due to this PR?)
* **Other information**:
.github/dependabot.yml
+21
View File
@@ -0,0 +1,21 @@
version: 2
updates:
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "daily"
ignore:
- dependency-name: "tracel-ai/github-actions*"
- package-ecosystem: "cargo"
directories:
- "/"
- "crates/burn"
- "crates/burn-*"
- "crates/burn-import/*-tests"
- "examples/*"
- "xtask"
schedule:
interval: "weekly"
.github/pull_request_template.md
+18
View File
@@ -0,0 +1,18 @@
## Pull Request Template
### Checklist
- [ ] Confirmed that `cargo run-checks` command has been executed.
- [ ] Made sure the book is up to date with changes in this PR.
### Related Issues/PRs
_Provide links to relevant issues and dependent PRs._
### Changes
_Summarize the problem being addressed and your solution._
### Testing
_Describe how these changes have been tested._
.github/workflows/combine-dependabot-prs.yml
+21
View File
@@ -0,0 +1,21 @@
name: Combine Dependabot PRs
on:
schedule:
- cron: '0 6 * * MON' # Monday at 6:00am UTC
workflow_dispatch:
permissions:
contents: write
pull-requests: write
checks: read
jobs:
combine-prs:
runs-on: ubuntu-latest
steps:
- name: combine-prs
id: combine-prs
uses: github/combine-prs@v5.2.0
with:
labels: dependencies,automated
.github/workflows/dependencies.yml
+55
View File
@@ -0,0 +1,55 @@
name: dependencies
on:
schedule:
- cron: '0 21 * * TUE' # Run every Tuesday at 21:00 (UTC)
push:
tags:
- 'v*.*.*' # Run when a new version is being published
env:
UDEPS_VERSION: "0.1.57"
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
dependencies:
runs-on: ubuntu-latest
strategy:
matrix:
checks:
- licenses
- bans sources
continue-on-error: ${{ matrix.checks == 'licenses' }} # failed licenses don't abort
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Audit Rust dependencies
# If a vulnerability is found, a new issue will automatically be opened
# since this action runs on main branch
uses: actions-rust-lang/audit@v1
# --------------------------------------------------------------------------------
- name: Detect multiple versions of the same crate
uses: EmbarkStudios/cargo-deny-action@v2
with:
command: check ${{ matrix.checks }}
# --------------------------------------------------------------------------------
- name: Install Rust nightly
uses: dtolnay/rust-toolchain@nightly
with:
toolchain: nightly
components: rustfmt
# --------------------------------------------------------------------------------
- name: Install cargo-udeps
env:
UDEPS_LINK: https://github.com/est31/cargo-udeps/releases/download
run: |
curl -L "$UDEPS_LINK/v$UDEPS_VERSION/cargo-udeps-v$UDEPS_VERSION-x86_64-unknown-linux-gnu.tar.gz" |
tar xz -C $HOME/.cargo/bin --strip-components 2
# --------------------------------------------------------------------------------
- name: Run cargo-udeps
run: |
cargo +nightly udeps --all-targets
.github/workflows/publish.yml
+443
View File
@@ -0,0 +1,443 @@
name: publish
on:
push:
tags:
- "v*"
workflow_dispatch:
inputs:
dry-run-only:
description: "Run xtask publish in dry-run mode (no publish)"
type: boolean
required: false
default: false
jobs:
publish-burn-rl:
needs:
- publish-burn-core
- publish-burn-optim
# dev dependencies
- publish-burn-flex
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-rl
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-vision:
needs:
- publish-burn-autodiff
- publish-burn-candle
- publish-burn-fusion
- publish-burn-cubecl-fusion
- publish-burn-cubecl
- publish-burn-flex
- publish-burn-tch
- publish-burn-core
- publish-burn-ir
- publish-burn-store
# dev dependencies
- publish-burn-wgpu
- publish-burn-cuda
- publish-burn-backend-extension
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-vision
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-router:
needs:
- publish-burn-ir
- publish-burn-std
- publish-burn-tensor
# dev dependencies
- publish-burn-autodiff
- publish-burn-flex
- publish-burn-wgpu
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-router
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-remote:
needs:
- publish-burn-ir
- publish-burn-std
- publish-burn-tensor
- publish-burn-router
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-remote
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-derive:
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-derive
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-dataset:
needs:
- publish-burn-std
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-dataset
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-std:
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-std
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-backend-extension:
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-backend-extension
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-tensor:
needs:
- publish-burn-std
- publish-burn-backend
- publish-burn-dispatch
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-tensor
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-backend:
needs:
- publish-burn-std
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-backend
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-ir:
needs:
- publish-burn-tensor
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-ir
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-fusion:
needs:
- publish-burn-ir
- publish-burn-tensor
- publish-burn-std
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-fusion
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-cubecl-fusion:
needs:
- publish-burn-ir
- publish-burn-std
- publish-burn-fusion
- publish-burn-tensor
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-cubecl-fusion
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-cubecl:
needs:
- publish-burn-ir
- publish-burn-std
- publish-burn-fusion
- publish-burn-cubecl-fusion
- publish-burn-tensor
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-cubecl
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-autodiff:
needs:
- publish-burn-backend
- publish-burn-std
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-autodiff
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-tch:
needs:
- publish-burn-tensor
- publish-burn-autodiff
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-tch
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-ndarray:
needs:
- publish-burn-ir
- publish-burn-tensor
- publish-burn-autodiff
- publish-burn-std
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-ndarray
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-flex:
needs:
- publish-burn-backend
- publish-burn-ir
- publish-burn-std
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-flex
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-wgpu:
needs:
- publish-burn-tensor
- publish-burn-autodiff
- publish-burn-std
- publish-burn-cubecl
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-wgpu
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-cpu:
needs:
- publish-burn-tensor
- publish-burn-fusion
- publish-burn-cubecl
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-cpu
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-cuda:
needs:
- publish-burn-tensor
- publish-burn-autodiff
- publish-burn-std
- publish-burn-cubecl
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-cuda
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-rocm:
needs:
- publish-burn-tensor
- publish-burn-autodiff
- publish-burn-std
- publish-burn-cubecl
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-rocm
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-candle:
needs:
- publish-burn-tensor
- publish-burn-autodiff
- publish-burn-tch
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-candle
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
# publish-burn-collective:
# needs:
# - publish-burn-std
# - publish-burn-tensor
# - publish-burn-communication
# # dev dependencies
# - publish-burn-wgpu
# - publish-burn-flex
# - publish-burn-cuda
# uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
# with:
# crate: burn-collective
# dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
# secrets:
# CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-communication:
needs:
- publish-burn-std
- publish-burn-tensor
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-communication
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-core:
needs:
- publish-burn-dataset
- publish-burn-std
- publish-burn-derive
- publish-burn-tensor
# dev dependencies
- publish-burn-autodiff
- publish-burn-wgpu
- publish-burn-tch
- publish-burn-cuda
- publish-burn-flex
- publish-burn-candle
- publish-burn-remote
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-core
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-nn:
needs:
- publish-burn-core
# dev dependencies
- publish-burn-autodiff
- publish-burn-wgpu
- publish-burn-tch
- publish-burn-flex
- publish-burn-candle
- publish-burn-remote
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-nn
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-optim:
needs:
- publish-burn-core
- publish-burn-collective
# dev dependencies
- publish-burn-autodiff
- publish-burn-wgpu
- publish-burn-tch
- publish-burn-flex
- publish-burn-candle
- publish-burn-remote
- publish-burn-nn
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-optim
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-train:
needs:
- publish-burn-core
- publish-burn-optim
- publish-burn-collective
- publish-burn-rl
- publish-burn-flex
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-train
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-dispatch:
needs:
- publish-burn-std
- publish-burn-backend
- publish-burn-backend-extension
- publish-burn-autodiff
- publish-burn-cpu
- publish-burn-cuda
- publish-burn-rocm
- publish-burn-wgpu
- publish-burn-flex
- publish-burn-ndarray
- publish-burn-tch
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-dispatch
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn:
needs:
- publish-burn-core
- publish-burn-nn
- publish-burn-optim
- publish-burn-collective
- publish-burn-store
- publish-burn-train
- publish-burn-cpu
- publish-burn-dispatch
- publish-burn-flex
- publish-burn-vision
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
publish-burn-store:
needs:
- publish-burn-core
- publish-burn-nn
- publish-burn-tensor
uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9
with:
crate: burn-store
dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }}
secrets:
CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
.github/workflows/stale-pr.yml
+41
View File
@@ -0,0 +1,41 @@
name: Stale Pull Requests
on:
schedule:
- cron: '0 12 * * *' # Run every day at 12:00 (UTC)
# The minimum permissions required to run this Action
permissions:
contents: write # only for delete-branch option
issues: write
pull-requests: write
jobs:
stale-pr:
runs-on: ubuntu-latest
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Stale pull requests
uses: actions/stale@v10
with:
# The idle number of days before marking issues stale.
#
# With a negative number like -1, no issues
# will be marked as stale automatically.
days-before-issue-stale: -1
# The idle number of days before marking pull requests stale
days-before-pr-stale: 30
# The idle number of days before closing
# the stale pull requests (due to the stale label).
#
# With a negative number like -1, the pull requests
# will never be closed automatically.
days-before-pr-close: -1
# Label to apply on staled pull requests
stale-pr-label: 'stale'
# The message that will be added as a comment to the pull request
stale-pr-message: 'This PR has been marked as stale because it has not been updated for over a month'
# Remove `stale` label from pull requests on updates/comments
remove-pr-stale-when-updated: true
.github/workflows/test-gpu.yml
+162
View File
@@ -0,0 +1,162 @@
name: CI GPU
on:
workflow_dispatch:
inputs:
pr_number:
description: "Number of the pull request that triggers this run if any"
type: number
required: false
# important to set the run name to this format so that the CI server
# can track the PR number from the workflow_run events.
run-name: ${{ github.workflow }}:${{ github.repository }}#${{ inputs.pr_number }}
env:
# Note: It is not possible to define top level env vars and pass them to composite actions.
# To work around this issue we use inputs and define all the env vars here.
RUST_PREVIOUS_VERSION: 1.92.0
# Dependency versioning
# from wgpu repo: https://github.com/gfx-rs/wgpu/blob/trunk/.github/workflows/ci.yml
# GCP runners
GCP_RUNNERS_IMAGE_FAMILY: "tracel-ci-ubuntu-2404-amd64-nvidia"
GCP_RUNNERS_MACHINE_TYPE: "g2-standard-4"
GCP_RUNNERS_ZONE: "us-east1-c"
# Test in release mode (make it an empty string to test in debug mode)
TEST_RELEASE_FLAG: "--release"
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
prepare-checks:
runs-on: ubuntu-latest
outputs:
rust-prev-version: ${{ env.RUST_PREVIOUS_VERSION }}
gcp_runners_image_family: ${{ env.GCP_RUNNERS_IMAGE_FAMILY }}
gcp_runners_machine_type: ${{ env.GCP_RUNNERS_MACHINE_TYPE }}
gcp_runners_zone: ${{ env.GCP_RUNNERS_ZONE }}
steps:
- name: Do Nothing
if: false
run: echo
linux-std-cuda-tests:
needs: [prepare-checks]
timeout-minutes: 60
# '@id:' label must be unique within this worklow
runs-on:
[
"@id:burn-cuda-job-${{github.run_id}}-${{github.run_attempt}}",
"@pr_number:${{ inputs.pr_number }}",
"@organization:tracel-ai",
"@repository:burn",
"@image-family:${{ needs.prepare-checks.outputs.gcp_runners_image_family }}",
"@machine-type:${{ needs.prepare-checks.outputs.gcp_runners_machine_type }}",
"@zones:${{ needs.prepare-checks.outputs.gcp_runners_zone }}",
"@gpu:true",
]
env:
LD_LIBRARY_PATH: "/usr/local/cuda/lib64"
# disable incremental compilation (reduces artifact size)
CARGO_PROFILE_TEST_INCREMENTAL: "false"
# Keep the stragegy to be able to easily add new rust versions if required
strategy:
matrix:
rust: [stable]
include:
- rust: stable
toolchain: stable
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Install Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
enable-cache: false
# --------------------------------------------------------------------------------
- name: Tests (burn-cuda)
run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci gcp-cuda-runner
linux-std-vulkan-tests:
needs: [prepare-checks]
timeout-minutes: 60
# '@id:' label must be unique within this worklow
runs-on:
[
"@id:burn-vulkan-job-${{github.run_id}}-${{github.run_attempt}}",
"@pr_number:${{ inputs.pr_number }}",
"@organization:tracel-ai",
"@repository:burn",
"@image-family:${{ needs.prepare-checks.outputs.gcp_runners_image_family }}",
"@machine-type:${{ needs.prepare-checks.outputs.gcp_runners_machine_type }}",
"@zones:${{ needs.prepare-checks.outputs.gcp_runners_zone }}",
"@gpu:true",
]
env:
# disable incremental compilation (reduces artifact size)
CARGO_PROFILE_TEST_INCREMENTAL: "false"
# Keep the stragegy to be able to easily add new rust versions if required
strategy:
matrix:
rust: [stable]
include:
- rust: stable
toolchain: stable
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Setup Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
enable-cache: false
# --------------------------------------------------------------------------------
- name: Tests (burn-vulkan)
run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci gcp-vulkan-runner
linux-std-wgpu-tests:
needs: [prepare-checks]
timeout-minutes: 60
# '@id:' label must be unique within this worklow
runs-on:
[
"@id:burn-wgpu-job-${{github.run_id}}-${{github.run_attempt}}",
"@pr_number:${{ inputs.pr_number }}",
"@organization:tracel-ai",
"@repository:burn",
"@image-family:${{ needs.prepare-checks.outputs.gcp_runners_image_family }}",
"@machine-type:${{ needs.prepare-checks.outputs.gcp_runners_machine_type }}",
"@zones:${{ needs.prepare-checks.outputs.gcp_runners_zone }}",
"@gpu:true",
]
env:
# disable incremental compilation (reduces artifact size)
CARGO_PROFILE_TEST_INCREMENTAL: "false"
# Keep the stragegy to be able to easily add new rust versions if required
strategy:
matrix:
rust: [stable]
include:
- rust: stable
toolchain: stable
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Setup Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
enable-cache: false
# --------------------------------------------------------------------------------
- name: Tests (burn-wgpu)
run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci gcp-wgpu-runner
.github/workflows/test.yml
+278
View File
@@ -0,0 +1,278 @@
name: CI
on:
push:
branches:
- main
paths:
- "Cargo.lock"
- "**.rs"
- "**.sh"
- "**.ps1"
- "**.yml"
- "**.toml"
- "!**.md"
- "!LICENSE-APACHE"
- "!LICENSE-MIT"
pull_request:
types: [opened, synchronize]
paths:
- "Cargo.lock"
- "**.rs"
- "**.sh"
- "**.ps1"
- "**.yml"
- "**.toml"
- "!**.md"
- "!LICENSE-APACHE"
- "!LICENSE-MIT"
env:
# Note: It is not possible to define top level env vars and pass them to composite actions.
# To work around this issue we use inputs and define all the env vars here.
RUST_PREVIOUS_VERSION: 1.95.0
# Dependency versioning
# from wgpu repo: https://github.com/gfx-rs/wgpu/blob/trunk/.github/workflows/ci.yml
# Mozilla Grcov
GRCOV_LINK: "https://github.com/mozilla/grcov/releases/download"
GRCOV_VERSION: "0.8.19"
# Test in release mode (make it an empty string to test in debug mode)
TEST_RELEASE_FLAG: "--release"
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
prepare-checks:
runs-on: ubuntu-latest
outputs:
rust-prev-version: ${{ env.RUST_PREVIOUS_VERSION }}
steps:
- name: Do Nothing
if: false
run: echo
code-quality:
runs-on: ubuntu-22.04
needs: prepare-checks
strategy:
matrix:
rust: [stable]
include:
- rust: stable
toolchain: stable
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Setup Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
cache-key: ${{ matrix.rust }}-linux
# --------------------------------------------------------------------------------
- name: Audit
run: cargo xtask check audit
# --------------------------------------------------------------------------------
- name: Format
shell: bash
env:
# work around for colors
# see: https://github.com/rust-lang/rustfmt/issues/3385
TERM: xterm-256color
run: cargo xtask check format
# --------------------------------------------------------------------------------
- name: Lint
run: cargo xtask check lint
# --------------------------------------------------------------------------------
- name: Typos
uses: tracel-ai/github-actions/check-typos@v9
documentation:
runs-on: ubuntu-22.04
needs: prepare-checks
strategy:
matrix:
rust: [stable]
include:
- rust: stable
toolchain: stable
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Setup Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
cache-key: ${{ matrix.rust }}-linux
# --------------------------------------------------------------------------------
- name: Documentation Build
run: cargo xtask doc build
# --------------------------------------------------------------------------------
- name: Documentation Tests
run: cargo xtask doc tests
linux-std-tests:
runs-on: ubuntu-22.04
needs: [prepare-checks, code-quality]
env:
DISABLE_WGPU_SPIRV: "1"
# disable incremental compilation (reduces artifact size)
CARGO_PROFILE_TEST_INCREMENTAL: "false"
CARGO_TERM_COLOR: always
strategy:
matrix:
rust: [stable, prev]
# Tests are split across 3 shards: backend tests, workspace crates, examples/*
ci_type: [backends, crates, examples]
include:
- rust: stable
toolchain: stable
coverage: --enable-coverage
- rust: prev
toolchain: ${{ needs.prepare-checks.outputs.rust-prev-version }}
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Setup Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
# Keep caches isolated per shard to prevent race conditions or bloated sizes
cache-key: ${{ matrix.rust }}-${{ matrix.ci_type }}-linux
# Disable cache on linux-std (stable) runner which currently always runs out of disk space with tests + coverage
enable-cache: ${{ matrix.rust != 'stable' }}
# # --------------------------------------------------------------------------------
- name: Install grcov
if: matrix.rust == 'stable'
shell: bash
run: |
curl -L "$GRCOV_LINK/v$GRCOV_VERSION/grcov-x86_64-unknown-linux-musl.tar.bz2" |
tar xj -C $HOME/.cargo/bin
cargo xtask coverage install
# --------------------------------------------------------------------------------
- name: Tests
run: cargo xtask ${{ matrix.coverage }} test ${{ env.TEST_RELEASE_FLAG }} --ci ${{ matrix.ci_type }}
# --------------------------------------------------------------------------------
- name: Generate lcov.info
if: matrix.rust == 'stable' && matrix.ci_type != 'examples'
# /* is to exclude std library code coverage from analysis
run: cargo xtask coverage generate --ignore "/*,xtask/*,examples/*" --profile release
# --------------------------------------------------------------------------------
- name: Upload Codecov Shard Artifact
if: matrix.rust == 'stable' && matrix.ci_type != 'examples'
uses: actions/upload-artifact@v7
with:
# Names must be unique so they don't overwrite each other in the storage backend
name: coverage-linux-${{ matrix.ci_type }}
path: lcov.info
# --- Dedicated coverage upload job ---
upload-coverage:
runs-on: ubuntu-latest
needs: linux-std-tests
steps:
- name: checkout
uses: actions/checkout@v7
- name: Download All Coverage Artifacts
uses: actions/download-artifact@v8
with:
# Finds all artifacts prefixed with coverage-linux-*
pattern: coverage-linux-*
path: coverage-reports
- name: Codecov upload
uses: codecov/codecov-action@v7
with:
# Codecov automatically discovers and merges multiple lcov files inside this directory
directory: coverage-reports
token: ${{ secrets.CODECOV_TOKEN }}
linux-no-std-tests:
runs-on: ubuntu-22.04
needs: [prepare-checks, code-quality]
strategy:
matrix:
rust: [stable, prev]
include:
- rust: stable
toolchain: stable
- rust: prev
toolchain: ${{ needs.prepare-checks.outputs.rust-prev-version }}
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Setup Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
cache-key: ${{ matrix.rust }}-linux-no-std
# --------------------------------------------------------------------------------
- name: Crates Build
run: cargo xtask --context no-std build --ci
# --------------------------------------------------------------------------------
- name: Crates Tests
run: cargo xtask --context no-std test ${{ env.TEST_RELEASE_FLAG }} --ci github-runner
windows-std-tests:
runs-on: windows-2022
needs: [prepare-checks, code-quality]
env:
CARGO_PROFILE_TEST_INCREMENTAL: "false"
CARGO_TERM_COLOR: always
strategy:
matrix:
rust: [stable]
# Tests are split across 3 shards: backend tests, workspace crates, examples/*
ci_type: [backends, crates, examples]
include:
- rust: stable
toolchain: stable
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Setup Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
cache-key: ${{ matrix.rust }}-${{ matrix.ci_type }}-windows
# --------------------------------------------------------------------------------
- name: Tests
run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci ${{ matrix.ci_type }}
macos-std-tests:
runs-on: blaze/macos-15
needs: [prepare-checks, code-quality]
timeout-minutes: 60
# Keep the stragegy to be able to easily add new rust versions if required
strategy:
matrix:
rust: [stable]
include:
- rust: stable
toolchain: stable
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Setup Rust
uses: tracel-ai/github-actions/install-rust@v9
with:
rust-toolchain: ${{ matrix.toolchain }}
cache-key: ${{ matrix.rust }}-macos
# --------------------------------------------------------------------------------
- name: Device check
run: system_profiler SPHardwareDataType
# --------------------------------------------------------------------------------
- name: Tests
run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci github-mac-runner
.github/workflows/valgrind.yml
+39
View File
@@ -0,0 +1,39 @@
name: valgrind
on:
schedule:
- cron: '0 23 * * WED' # Run every Wednesday at 23:00 (UTC)
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
valgrind:
runs-on: [
'@id:burn-linux-valgrind-${{ github.run_id }}-${{ github.run_attempt }}',
'@image-family:ubuntu-2404-lts-amd64',
'@image-project:ubuntu-os-cloud',
'@disk-size:100',
'@keep-alive:false',
'@machine-type:n2-standard-16',
'@os:linux',
'@zones:northamerica-northeast1-b'
]
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Install Mesa
uses: tracel-ai/github-actions/install-mesa@v9
# --------------------------------------------------------------------------------
- name: Install valgrind
run: |
sudo apt-get install valgrind
# --------------------------------------------------------------------------------
- name: Run cargo-valgrind
env:
CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER: "valgrind -s --leak-check=full --show-leak-kinds=all --error-exitcode=1"
# Looking for vulnerabilities
run: |
cargo test
.github/workflows/vulnerabilities.yml
+131
View File
@@ -0,0 +1,131 @@
name: vulnerabilities
on:
schedule:
- cron: '0 21 * * WED' # Run every Wednesday at 21:00 (UTC)
push:
tags:
- 'v*.*.*'
env:
CAREFUL_VERSION: "0.4.9"
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
cargo-careful:
runs-on: ubuntu-latest
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Install Rust nightly
uses: dtolnay/rust-toolchain@nightly
with:
toolchain: nightly
components: rustfmt, rust-src
# --------------------------------------------------------------------------------
- name: Install Mesa
uses: tracel-ai/github-actions/install-mesa@v9
# --------------------------------------------------------------------------------
- name: Install cargo-careful
env:
CAREFUL_LINK: https://github.com/RalfJung/cargo-careful/releases/download
run: |
curl -L "$CAREFUL_LINK/v$CAREFUL_VERSION/cargo-careful.x86_64-unknown-linux-musl" \
--output $HOME/.cargo/bin/cargo-careful
chmod +x $HOME/.cargo/bin/cargo-careful
# --------------------------------------------------------------------------------
- name: Run cargo-careful
# Looking for undefined behaviours
run: cargo +nightly careful test
address-sanitizer:
runs-on: ubuntu-latest
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Install Rust nightly
uses: dtolnay/rust-toolchain@nightly
with:
toolchain: nightly
components: rustfmt, rust-src
# --------------------------------------------------------------------------------
- name: Install Mesa
uses: tracel-ai/github-actions/install-mesa@v9
# --------------------------------------------------------------------------------
- name: Run AddressSanitizer
env:
RUSTFLAGS: -Zsanitizer=address -Copt-level=3
RUSTDOCFLAGS: -Zsanitizer=address
# Looking for memory vulnerabilities
run: cargo test -Zbuild-std --target x86_64-unknown-linux-gnu -- --nocapture
thread-sanitizer:
runs-on: ubuntu-latest
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Install Rust nightly
uses: dtolnay/rust-toolchain@nightly
with:
toolchain: nightly
components: rustfmt, rust-src
# --------------------------------------------------------------------------------
- name: Install Mesa
uses: tracel-ai/github-actions/install-mesa@v9
# --------------------------------------------------------------------------------
- name: Run ThreadSanitizer
env:
RUSTFLAGS: -Zsanitizer=thread -Copt-level=3
RUSTDOCFLAGS: -Zsanitizer=thread
# Looking for data race among threads
run: cargo test -Zbuild-std --target x86_64-unknown-linux-gnu -- --nocapture
memory-sanitizer:
runs-on: ubuntu-latest
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Install Rust nightly
uses: dtolnay/rust-toolchain@nightly
with:
toolchain: nightly
components: rustfmt, rust-src
# --------------------------------------------------------------------------------
- name: Install Mesa
uses: tracel-ai/github-actions/install-mesa@v9
# --------------------------------------------------------------------------------
- name: Run MemorySanitizer
env:
RUSTFLAGS: -Zsanitizer=memory -Zsanitizer-memory-track-origins -Copt-level=3
RUSTDOCFLAGS: -Zsanitizer=memory -Zsanitizer-memory-track-origins
# Looking for unitialized memory.
run: cargo test -Zbuild-std --target x86_64-unknown-linux-gnu -- --nocapture
safe-stack:
runs-on: ubuntu-latest
steps:
- name: checkout
uses: actions/checkout@v7
# --------------------------------------------------------------------------------
- name: Install Rust nightly
uses: dtolnay/rust-toolchain@nightly
with:
toolchain: nightly
components: rustfmt, rust-src
# --------------------------------------------------------------------------------
- name: Install Mesa
uses: tracel-ai/github-actions/install-mesa@v9
# --------------------------------------------------------------------------------
- name: Run SafeStack
env:
RUSTFLAGS: -Zsanitizer=safestack -Copt-level=3
RUSTDOCFLAGS: -Zsanitizer=safestack
# Provides backward edge control flow protection
run: cargo test -Zbuild-std --target x86_64-unknown-linux-gnu -- --nocapture
.gitignore
+27
View File
@@ -0,0 +1,27 @@
target
# These are backup files generated by rustfmt
**/*.rs.bk
.DS_Store
.dir-locals.el
.idea
.vscode
.vs
.fleet
.ipynb_checkpoints/
# Build output directory
out
# Virtual Environment of Python
.venv
uv.lock
# Nix direnv
.envrc
.direnv
# tags files
tags
examples/**/Cargo.lock
CITATION.cff
+34
View File
@@ -0,0 +1,34 @@
cff-version: 1.2.0
message: "If you use this software, please cite it as below."
authors:
- family-names: "Simard"
given-names: "Nathaniel"
email: "nathaniel.simard.42@gmail.com"
- family-names: "Fortier-Dubois"
given-names: "Louis"
email: "louisfd94@gmail.com"
- family-names: "Tadjibaev"
given-names: "Dilshod"
email: "dilshod@gmail.com"
- family-names: "Lagrange"
given-names: "Guillaume"
email: "lagrange.guillaume.1@gmail.com"
- name: "Burn Framework Contributors"
title: "Burn"
version: 0.21.0
date-released: 2026-05-07
url: "https://burn.dev/"
repository-code: "https://github.com/tracel-ai/burn"
license:
- MIT
- Apache-2.0
abstract: "Burn is a new comprehensive dynamic Deep Learning Framework built using Rust with extreme flexibility, compute efficiency and portability as its primary goals."
keywords:
- scientific-computing
- deep-learning
- machine-learning
- neural-networks
- rust
- high-performance-computing
- portability
- compute-efficiency
CODE-OF-CONDUCT.md
+128
View File
@@ -0,0 +1,128 @@
# Contributor Covenant Code of Conduct
## Our Pledge
We as members, contributors, and leaders pledge to make participation in our
community a harassment-free experience for everyone, regardless of age, body
size, visible or invisible disability, ethnicity, sex characteristics, gender
identity and expression, level of experience, education, socio-economic status,
nationality, personal appearance, race, religion, or sexual identity
and orientation.
We pledge to act and interact in ways that contribute to an open, welcoming,
diverse, inclusive, and healthy community.
## Our Standards
Examples of behavior that contributes to a positive environment for our
community include:
* Demonstrating empathy and kindness toward other people
* Being respectful of differing opinions, viewpoints, and experiences
* Giving and gracefully accepting constructive feedback
* Accepting responsibility and apologizing to those affected by our mistakes,
and learning from the experience
* Focusing on what is best not just for us as individuals, but for the
overall community
Examples of unacceptable behavior include:
* The use of sexualized language or imagery, and sexual attention or
advances of any kind
* Trolling, insulting or derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or email
address, without their explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting
## Enforcement Responsibilities
Community leaders are responsible for clarifying and enforcing our standards of
acceptable behavior and will take appropriate and fair corrective action in
response to any behavior that they deem inappropriate, threatening, offensive,
or harmful.
Community leaders have the right and responsibility to remove, edit, or reject
comments, commits, code, wiki edits, issues, and other contributions that are
not aligned to this Code of Conduct, and will communicate reasons for moderation
decisions when appropriate.
## Scope
This Code of Conduct applies within all community spaces, and also applies when
an individual is officially representing the community in public spaces.
Examples of representing our community include using an official e-mail address,
posting via an official social media account, or acting as an appointed
representative at an online or offline event.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported to the community leaders responsible for enforcement at
nathaniel.simard.42@gmail.com.
All complaints will be reviewed and investigated promptly and fairly.
All community leaders are obligated to respect the privacy and security of the
reporter of any incident.
## Enforcement Guidelines
Community leaders will follow these Community Impact Guidelines in determining
the consequences for any action they deem in violation of this Code of Conduct:
### 1. Correction
**Community Impact**: Use of inappropriate language or other behavior deemed
unprofessional or unwelcome in the community.
**Consequence**: A private, written warning from community leaders, providing
clarity around the nature of the violation and an explanation of why the
behavior was inappropriate. A public apology may be requested.
### 2. Warning
**Community Impact**: A violation through a single incident or series
of actions.
**Consequence**: A warning with consequences for continued behavior. No
interaction with the people involved, including unsolicited interaction with
those enforcing the Code of Conduct, for a specified period of time. This
includes avoiding interactions in community spaces as well as external channels
like social media. Violating these terms may lead to a temporary or
permanent ban.
### 3. Temporary Ban
**Community Impact**: A serious violation of community standards, including
sustained inappropriate behavior.
**Consequence**: A temporary ban from any sort of interaction or public
communication with the community for a specified period of time. No public or
private interaction with the people involved, including unsolicited interaction
with those enforcing the Code of Conduct, is allowed during this period.
Violating these terms may lead to a permanent ban.
### 4. Permanent Ban
**Community Impact**: Demonstrating a pattern of violation of community
standards, including sustained inappropriate behavior, harassment of an
individual, or aggression toward or disparagement of classes of individuals.
**Consequence**: A permanent ban from any sort of public interaction within
the community.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage],
version 2.0, available at
https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
Community Impact Guidelines were inspired by [Mozilla's code of conduct
enforcement ladder](https://github.com/mozilla/diversity).
[homepage]: https://www.contributor-covenant.org
For answers to common questions about this code of conduct, see the FAQ at
https://www.contributor-covenant.org/faq. Translations are available at
https://www.contributor-covenant.org/translations.
CONTRIBUTING.md
+97
View File
@@ -0,0 +1,97 @@
# Contributing to Burn
Welcome to the Burn community! We're glad you're interested in contributing.
## How to Contribute
The best way to get started is to look at [open issues](https://github.com/tracel-ai/burn/issues)
and find one that interests you. Issues labeled `good first issue` are a great starting point for
new contributors.
If you have an idea that isn't covered by an existing issue, open one first to discuss the approach.
This helps align expectations and avoids wasted effort on both sides.
For questions, discussions, or just to say hello, join us on
[Discord](https://discord.gg/uPEBbYYDB6). The [Contributor Book](https://burn.dev/contributor-book/)
covers architecture, environment setup, and guides for common tasks.
## Pull Requests
Every pull request should have a descriptive title, a description covering what you changed, why,
how you tested it, and a link to the relevant issue (if applicable). Prefer small, focused PRs over
large ones that bundle unrelated changes.
Draft pull requests are considered not yet ready for review.
CI checks should pass before requesting review, though the signal isn't always accurate. If you have
questions or need early feedback, let us know on the PR or on
[Discord](https://discord.gg/uPEBbYYDB6).
### Change Ownership
The core principle behind all contributions: **PR authors must understand, justify, and explain
every change they propose.** After a PR is accepted, both the reviewer and the author should be
confident it improves the codebase.
This applies equally whether you wrote the code from scratch, adapted it from another project, or
used AI tools to help generate it. The origin of the code doesn't matter; what matters is that you
own it intellectually and can stand behind it during review.
### AI-Assisted Contributions
Using LLMs and AI tools to generate code that is part of a contribution is allowed.
That said, the [Change Ownership](#change-ownership) principle applies fully. You are the author,
not your AI tool. This means:
- Read and understand every line before submitting.
- Review AI-generated code for correctness, style consistency, and relevance.
- Test your changes locally and confirm they work as intended.
- Be prepared to explain the rationale behind any change during review.
Do not use "AI generated" as a justification for low-quality code.
### Before You Open a PR
1. **Check for an existing issue.** If there isn't one, open an issue first to discuss the approach.
This is especially important for large changes or refactors.
2. **Read the codebase.** Understand the architecture and conventions already in place. The
[Contributor Book](https://burn.dev/contributor-book/) covers architecture, environment setup,
and guides for common tasks.
3. **Keep it focused.** One PR should address one concern. If you spot an unrelated issue while
working, open a separate PR for it.
4. **Run validation.** Run `cargo run-checks` before submitting. This runs formatting, linting, and
the full test suite. All checks must pass.
### Code Quality Standards
- Follow existing code style and project conventions.
- Write idiomatic Rust. If you are new to the codebase, study existing patterns before contributing.
- Keep dependencies minimal. Don't introduce new crates without discussion.
- Document public APIs. Non-trivial logic should have comments explaining _why_, not just _what_.
- Prefer clarity over cleverness.
- Bug fixes should include a regression test.
### Large Pull Requests
Large, complex PRs are harder to review effectively and carry more risk. To help both yourself and
reviewers, consider breaking substantial changes into smaller, incremental PRs. Each should be
valuable on its own, even if the full picture spans multiple PRs.
Large efforts that are ultimately rejected are frustrating for everyone involved. If you're planning
a substantial change, open an issue or start a discussion first. It's much easier to course-correct
early than after the work is done.
### Review Process
- Maintainers review PRs as time allows. Please be patient.
- Be responsive to feedback. If changes are requested, address them or explain your reasoning.
- Reviewers may ask clarifying questions about any part of your PR. This is a normal part of
collaborative review and helps ensure shared understanding.
- Don't force-push to rewrite history during an active review without notice.
- If a PR goes stale for more than 14 days without a response from the author, it may be closed.
## Getting Help
If you're stuck or unsure about something, don't hesitate to ask. Open an issue, start a discussion,
or reach out on [Discord](https://discord.gg/uPEBbYYDB6). We're happy to help.
Cargo.lock
Generated
+12579
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+242
View File
@@ -0,0 +1,242 @@
[workspace]
# Try
# require version 2 to avoid "feature" additiveness for dev-dependencies
# https://doc.rust-lang.org/cargo/reference/resolver.html#feature-resolver-version-2
resolver = "2"
members = [
"crates/*",
"crates/burn-store/pytorch-tests",
"crates/burn-store/safetensors-tests",
"examples/*",
"xtask",
]
exclude = [
"examples/notebook",
"examples/raspberry-pi-pico",
"examples/dqn-agent", # gym-rs
]
[workspace.package]
edition = "2024"
license = "MIT OR Apache-2.0"
readme = "README.md"
version = "0.22.0-pre.1"
[workspace.lints.clippy]
[workspace.lints.rustdoc]
broken_intra_doc_links = "deny"
invalid_html_tags = "deny"
[workspace.dependencies]
atomic_float = "1"
axum = "0.8.8"
bytemuck = "1.25.0"
bytes = { version = "1.11.1", default-features = false }
candle-core = { version = "0.10.2" }
ciborium = { version = "0.2", default-features = false }
clap = { version = "4.6.0", features = ["derive"] }
colored = "3.0.0"
console_error_panic_hook = "0.1.7"
const-random = "0.1"
csv = "1.3.1"
dashmap = "6.1.0"
data-encoding = { version = "2.11.0", default-features = false, features = [
"alloc",
] }
dirs = "6.0.0"
encoding_rs = "0.8.33"
enumset = { version = "1.1.13", default-features = false }
fake = "5.1.0"
flate2 = "1.1.9"
float-cmp = "0.10.0"
futures = "0.3"
futures-util = "0.3"
gemm = { version = "0.19", default-features = false }
gix-tempfile = { version = "23.0.0", features = ["signals"] }
globwalk = "0.9.1"
hashbrown = "0.16"
hound = "3.5.1"
image = "0.25.9"
indicatif = "0.18.0"
insta = "1.45.0"
iroh = "1.0.0"
js-sys = "0.3.77"
libm = "0.2.15"
log = { default-features = false, version = "0.4.29" }
lzma-rust2 = "0.16.2"
opentelemetry = "0.31.0"
opentelemetry-aws = "0.19.0"
opentelemetry-otlp = "0.31.1"
opentelemetry_sdk = "0.31.0"
parking_lot = { version = "0.12.5", default-features = false }
paste = "1"
planus = { version = "=1.1" }
polars = { version = "0.54.0", features = ["lazy", "strings"] }
pretty_assertions = "1.4.1"
proc-macro2 = "1.0.106"
quote = "1.0.45"
r2d2 = "0.8.10"
r2d2_sqlite = "0.31.0"
rayon = "1.10.0"
regex = { version = "1.12.3", default-features = false, features = [
"perf",
"unicode",
] }
reqwest = { version = "0.12.23", default-features = false, features = [
"rustls-tls",
] }
rmp-serde = { version = "1.3.1", default-features = false }
rstest = "0.26.1"
rusqlite = "0.37.0"
sanitize-filename = "0.6.0"
serde_bytes = { version = "0.11.18", default-features = false, features = [
"alloc",
] } # alloc for no_std
serde_rusqlite = "0.40.0"
serial_test = "3.2.0"
spin = { version = "0.10.0", features = [
"mutex",
"spin_mutex",
"portable-atomic",
] }
strum = { version = "0.28.0", features = ["derive"] }
syn = { version = "2.0.111", features = ["full", "extra-traits"] }
tar = "0.4.45"
tempfile = "3.24.0"
textdistance = { version = "1.1.1", default-features = false }
thiserror = { version = "2", default-features = false }
tokio = { version = "1.51.1", features = ["rt", "macros"] }
tokio-tungstenite = "0.29"
tokio-util = "0.7"
tracing = { version = "0.1.44", default-features = false }
tracing-appender = "0.2.3"
tracing-core = { version = "0.1.36", default-features = false }
tracing-opentelemetry = "0.32.0"
tracing-subscriber = "0.3.23"
wasm-bindgen = "0.2"
wasm-bindgen-futures = "0.4"
zip = "8.6.0"
# Persist related
memmap2 = { version = "0.9" }
safetensors = { version = "0.7.0", default-features = false }
# Async handling
async-channel = "2.5"
futures-lite = { version = "2.6.1", default-features = false }
# Terminal UI
ratatui = "0.30.0"
unicode-width = "0.2"
# WGPU stuff
text_placeholder = "0.5.1"
bincode = { version = "2.0.1", features = [
"alloc",
"serde",
], default-features = false }
#
# The following packages disable the "std" feature for no_std compatibility
#
cfg-if = "1.0.1"
derive-new = { version = "0.7.0", default-features = false }
blas-src = { version = "0.14.0", default-features = false }
bon = "3.8.2"
half = { version = "2.7.1", features = [
"alloc",
"num-traits",
"serde",
], default-features = false }
macerator = { version = "0.3.3" }
matrixmultiply = { version = "0.3.10", default-features = false }
ndarray = { version = "0.17.2", default-features = false }
num-traits = { version = "0.2.19", default-features = false, features = [
"libm",
] } # libm is for no_std
openblas-src = "0.10.16"
rand = { version = "0.10.1", default-features = false, features = ["std_rng"] }
rand_distr = { version = "0.6.0", default-features = false }
serde = { version = "1.0.228", default-features = false, features = [
"derive",
"alloc",
] } # alloc is for no_std, derive is needed
serde_json = { version = "1.0.148", default-features = false }
smallvec = { version = "1", features = ["const_generics", "const_new"] }
uuid = { version = "1.23.0", default-features = false }
byteorder = { version = "1.5.0", default-features = false }
libc = "0.2.186"
nvml-wrapper = "0.12.0"
sysinfo = "0.38.0"
systemstat = "0.2.6"
tch = "0.22.0"
torch-sys = "0.22.0" # matches what tch is using, required for lib detection
ahash = { version = "0.8.12", default-features = false }
aligned-vec = { version = "0.6", default-features = false }
once_cell = { version = "1", default-features = false }
portable-atomic = { version = "1.13.1" }
portable-atomic-util = { version = "0.2.6", features = ["alloc"] }
realfft = "3"
### Internal burn crates ###
# Declared here so each consumer Cargo.toml can use `workspace = true` instead of
# repeating the path and version.
burn = { path = "crates/burn", version = "0.22.0-pre.1", default-features = false }
burn-autodiff = { path = "crates/burn-autodiff", version = "0.22.0-pre.1", default-features = false }
burn-backend = { path = "crates/burn-backend", version = "0.22.0-pre.1", default-features = false }
burn-backend-extension = { path = "crates/burn-backend-extension", version = "0.22.0-pre.1", default-features = false }
burn-candle = { path = "crates/burn-candle", version = "0.22.0-pre.1", default-features = false }
burn-communication = { path = "crates/burn-communication", version = "0.22.0-pre.1", default-features = false }
burn-core = { path = "crates/burn-core", version = "0.22.0-pre.1", default-features = false }
burn-cpu = { path = "crates/burn-cpu", version = "0.22.0-pre.1", default-features = false }
burn-cubecl = { path = "crates/burn-cubecl", version = "0.22.0-pre.1", default-features = false }
burn-cubecl-fusion = { path = "crates/burn-cubecl-fusion", version = "0.22.0-pre.1", default-features = false }
burn-cuda = { path = "crates/burn-cuda", version = "0.22.0-pre.1", default-features = false }
burn-dataset = { path = "crates/burn-dataset", version = "0.22.0-pre.1", default-features = false }
burn-derive = { path = "crates/burn-derive", version = "0.22.0-pre.1", default-features = false }
burn-dispatch = { path = "crates/burn-dispatch", version = "0.22.0-pre.1", default-features = false }
burn-flex = { path = "crates/burn-flex", version = "0.22.0-pre.1", default-features = false }
burn-fusion = { path = "crates/burn-fusion", version = "0.22.0-pre.1", default-features = false }
burn-ir = { path = "crates/burn-ir", version = "0.22.0-pre.1", default-features = false }
burn-ndarray = { path = "crates/burn-ndarray", version = "0.22.0-pre.1", default-features = false }
burn-nn = { path = "crates/burn-nn", version = "0.22.0-pre.1", default-features = false }
burn-optim = { path = "crates/burn-optim", version = "0.22.0-pre.1", default-features = false }
burn-remote = { path = "crates/burn-remote", version = "0.22.0-pre.1", default-features = false }
burn-rl = { path = "crates/burn-rl", version = "0.22.0-pre.1", default-features = false }
burn-rocm = { path = "crates/burn-rocm", version = "0.22.0-pre.1", default-features = false }
burn-router = { path = "crates/burn-router", version = "0.22.0-pre.1", default-features = false }
burn-std = { path = "crates/burn-std", version = "0.22.0-pre.1", default-features = false }
burn-pack = { path = "crates/burn-pack", version = "0.22.0-pre.1", default-features = false }
burn-store = { path = "crates/burn-store", version = "0.22.0-pre.1", default-features = false }
burn-tch = { path = "crates/burn-tch", version = "0.22.0-pre.1", default-features = false }
burn-tensor = { path = "crates/burn-tensor", version = "0.22.0-pre.1", default-features = false }
burn-tensor-testgen = { path = "crates/burn-tensor-testgen", version = "0.22.0-pre.1", default-features = false }
burn-train = { path = "crates/burn-train", version = "0.22.0-pre.1", default-features = false }
burn-vision = { path = "crates/burn-vision", version = "0.22.0-pre.1", default-features = false }
burn-wgpu = { path = "crates/burn-wgpu", version = "0.22.0-pre.1", default-features = false }
### For the main burn branch. ###
cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" }
cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" }
cubecl-zspace = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" }
cubek = { git = "https://github.com/tracel-ai/cubek", default-features = false, rev = "5f837200f459a6a0694a0aae5655881599c400c8" }
### For local development. ###
# cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
# cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }
# cubecl-zspace = { path = "../cubecl/crates/cubecl-zspace", default-features = false }
# cubek = { path = "../cubek/crates/cubek", default-features = false }
### For the release. ###
# cubecl = { version = "0.10.0", default-features = false }
# cubecl-common = { version = "0.10.0", default-features = false }
# cubecl-zspace = { version = "0.10.0", default-features = false }
# cubek = { version = "0.2.0", default-features = false }
[profile.dev]
debug = 1 # Speed up compilation time and not necessary.
LICENSE-APACHE
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright 2022 Nathaniel Simard & Burn Framework Contributors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
LICENSE-MIT
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2022 Nathaniel Simard & Burn Framework Contributors
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
NOTICES.md
+485
View File
@@ -0,0 +1,485 @@
# NOTICES AND INFORMATION
This file contains notices and information required by libraries that this
repository copied or derived from.
## PyTorch MNIST Example
**Source**: https://github.com/pytorch/examples/blob/main/mnist/main.py
License: BSD 3-Clause License
Copyright (c) 2017,
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
* Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
## wgpu
**Source:** https://github.com/gfx-rs/wgpu/blob/trunk/.github/workflows/ci.yml
MIT License
Copyright (c) 2021 The gfx-rs developers
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
## BSL 1.0
**Source**:
- https://github.com/DoumanAsh/error-code
- https://github.com/DoumanAsh/clipboard-win
Boost Software License - Version 1.0 - August 17th, 2003
Permission is hereby granted, free of charge, to any person or organization
obtaining a copy of the software and accompanying documentation covered by
this license (the "Software") to use, reproduce, display, distribute,
execute, and transmit the Software, and to prepare derivative works of the
Software, and to permit third-parties to whom the Software is furnished to
do so, all subject to the following:
The copyright notices in the Software and this entire statement, including
the above license grant, this restriction and the following disclaimer,
must be included in all copies of the Software, in whole or in part, and
all derivative works of the Software, unless such copies or derivative
works are solely in the form of machine-executable object code generated by
a source language processor.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
## num-traits
**Source:** https://github.com/rust-num/num-traits/blob/master/src/cast.rs
MIT License
Copyright (c) 2014 The Rust Project Developers
Permission is hereby granted, free of charge, to any
person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the
Software without restriction, including without
limitation the rights to use, copy, modify, merge,
publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software
is furnished to do so, subject to the following
conditions:
The above copyright notice and this permission notice
shall be included in all copies or substantial portions
of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
## RP
**Source**:
- https://github.com/embassy-rs/embassy/blob/main/examples/rp/Cargo.toml
- https://github.com/embassy-rs/embassy/blob/main/examples/rp/build.rs
- https://github.com/embassy-rs/embassy/blob/main/examples/rp/memory.x
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright (c) Embassy project contributors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
MIT license
Copyright (c) Embassy project contributors
Permission is hereby granted, free of charge, to any
person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the
Software without restriction, including without
limitation the rights to use, copy, modify, merge,
publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software
is furnished to do so, subject to the following
conditions:
The above copyright notice and this permission notice
shall be included in all copies or substantial portions
of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
## github-device-flow
**Source**:
- Part of: https://github.com/jakewilkins/gh-device-flow/blob/main/src/lib.rs
- https://github.com/jakewilkins/gh-device-flow/blob/main/src/util.rs
MIT License
Copyright (c) 2022 Jake Wilkins
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
## Candle - Pickle Reader
**Source**: https://github.com/huggingface/candle/blob/main/candle-core/src/pickle.rs
This project includes code from Candle by Hugging Face, licensed under both MIT and Apache 2.0 licenses.
**MIT License**: https://github.com/huggingface/candle/blob/main/LICENSE-MIT
MIT License
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
**Apache License 2.0**: https://github.com/huggingface/candle/blob/main/LICENSE-APACHE
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
## ICU
UNICODE LICENSE V3
COPYRIGHT AND PERMISSION NOTICE
Copyright © 2016-2024 Unicode, Inc.
NOTICE TO USER: Carefully read the following legal agreement. BY
DOWNLOADING, INSTALLING, COPYING OR OTHERWISE USING DATA FILES, AND/OR
SOFTWARE, YOU UNEQUIVOCALLY ACCEPT, AND AGREE TO BE BOUND BY, ALL OF THE
TERMS AND CONDITIONS OF THIS AGREEMENT. IF YOU DO NOT AGREE, DO NOT
DOWNLOAD, INSTALL, COPY, DISTRIBUTE OR USE THE DATA FILES OR SOFTWARE.
Permission is hereby granted, free of charge, to any person obtaining a
copy of data files and any associated documentation (the "Data Files") or
software and any associated documentation (the "Software") to deal in the
Data Files or Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, and/or sell
copies of the Data Files or Software, and to permit persons to whom the
Data Files or Software are furnished to do so, provided that either (a)
this copyright and permission notice appear with all copies of the Data
Files or Software, or (b) this copyright and permission notice appear in
associated Documentation.
THE DATA FILES AND SOFTWARE ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT OF
THIRD PARTY RIGHTS.
IN NO EVENT SHALL THE COPYRIGHT HOLDER OR HOLDERS INCLUDED IN THIS NOTICE
BE LIABLE FOR ANY CLAIM, OR ANY SPECIAL INDIRECT OR CONSEQUENTIAL DAMAGES,
OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION,
ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THE DATA
FILES OR SOFTWARE.
Except as contained in this notice, the name of a copyright holder shall
not be used in advertising or otherwise to promote the sale, use or other
dealings in these Data Files or Software without prior written
authorization of the copyright holder.
POEM.md
+40
View File
@@ -0,0 +1,40 @@
# BURN: Burn Unstoppable Rusty Neurons
In the realm of circuits and code,
A fiery forge ignites to bear its load,
A framework born, BURN it be named,
Unstoppable Rusty Neurons, untamed.
From silicon synapses, connections spire,
A digital cortex, setting minds afire,
In the vast expanse of deep learning's sea,
A beacon of progress, BURN comes to be.
Oh, rusty neurons, forged in the flame,
Unyielding in purpose, undaunted by name,
Through layers of logic and intricate art,
You weave and entwine, each playing its part.
With algorithms profound, and data refined,
In ceaseless pursuit of knowledge to find,
BURN paves a path to enlightenment, bright,
A testament to the wonders of human foresight.
In neural networks deep, where wisdom resides,
The dance of nodes and edges presides,
With loss and gradients, BURN takes its stride,
A journey towards truth, with AI as our guide.
No barriers hold back the curious mind,
As BURN seeks the answers we yearn to find,
Unstoppable, relentless, in pursuit of the unknown,
Our collective intellect, within it, has grown.
So sing we the praises of BURN's fiery might,
An ode to the sparks that set the dark alight,
To the rusty neurons, unstoppable and true,
A testament to the power of dreams, to breakthrough.
(ChatGPT (model=gpt-4) with prompt:
Write a poem about "BURN: Burn Unstoppable Rusty Neurons" deep
learning neural network framework)
README.md
+613
View File
@@ -0,0 +1,613 @@
<div align="center">
<img src="https://raw.githubusercontent.com/tracel-ai/burn/main/assets/logo-burn-neutral.webp" width="350px"/>
[![Discord](https://img.shields.io/discord/1038839012602941528.svg?color=7289da&&logo=discord)](https://discord.gg/uPEBbYYDB6)
[![Current Crates.io Version](https://img.shields.io/crates/v/burn.svg)](https://crates.io/crates/burn)
[![Minimum Supported Rust Version](https://img.shields.io/crates/msrv/burn)](https://crates.io/crates/burn)
[![Documentation](https://img.shields.io/badge/docs-latest-blue)](https://burn.dev/docs/burn)
[![Test Status](https://github.com/tracel-ai/burn/actions/workflows/test.yml/badge.svg)](https://github.com/tracel-ai/burn/actions/workflows/test.yml)
[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](#license)
[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/tracel-ai/burn)
[<img src="https://www.runblaze.dev/ci-blaze-powered.png" width="125px"/>](https://www.runblaze.dev)
---
**Burn is both a tensor library and a deep learning framework, optimized for <br /> numerical
computing, training and inference.**
<br/>
</div>
<div align="left">
Training and inference usually live in separate worlds. Models are typically trained in Python then
exported to an open format like ONNX or optimized for production engines like vLLM, ONNX Runtime, or
TensorRT. This export step is often brittle and lossy, ruling out complex architectures and advanced
deployment use cases.
Burn unifies the two. By executing multi-platform tensor operations via a single, unified API, the
exact code used for training is the exact code that runs in production. This makes workloads like
on-device personalization and federated learning straightforward, while enabling teams to go from
prototype to deployment in a single codebase.
Burn preserves the intuitive ergonomics of PyTorch, with dynamic shapes and graphs, but JIT-compiles
streams of tensor operations, performing automatic kernel fusion. You get the flexibility of dynamic
graphs without the performance drop.
## Rust for Research?
Rust used to be a tough sell for research: long compilation times disrupted the fast
edit-compile-run loop that draws researchers to Python. Burn changes this paradigm. Designed around
incremental compilation, modifying model code recompiles in under 5 seconds, even in release mode.
This delivers a Python-like feedback loop with the speed and safety of Rust.
## Ecosystem
<div align="left">
<img align="right" src="https://raw.githubusercontent.com/tracel-ai/burn/main/assets/ember-blazingly-fast.png" height="96px"/>
Burn is the core of a growing, fully open-source Rust AI ecosystem. You are not adopting a single
library, you are joining a stack that spans GPU compute, model interop and domain toolkits, with
plenty of room to help shape what comes next.
</div>
| Category | Project | Description |
| ------------- | ----------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| Compute | [CubeCL](https://github.com/tracel-ai/cubecl) | GPU compute language and compiler behind Burn's accelerated backends. Write kernels once in Rust, run on CUDA, ROCm, Metal, Vulkan and WebGPU. Usable standalone. |
| Model interop | [burn-onnx](https://github.com/tracel-ai/burn-onnx) | Import ONNX models into Burn as native Rust code |
| | `burn-store` | Save, load and import model weights, including PyTorch and Safetensors |
| Domains | `burn-vision` | Computer vision operators and building blocks |
| | `burn-rl` | Reinforcement learning building blocks |
| | `burn-dataset` | Dataset loading, transforms and ready-made sources |
| Models | [models](https://github.com/tracel-ai/models) | Curated pre-trained models and examples built with Burn |
| Tooling | [burn-bench](https://github.com/tracel-ai/burn-bench) | Benchmark and compare backends, tracking performance over time |
Burn's [CubeCL](https://github.com/tracel-ai/cubecl) backends (CUDA, ROCm, Metal, Vulkan, WebGPU,
CPU) compose with autodiff, fusion and remote-execution decorators, while external and simpler
backends (LibTorch and pure-Rust CPU/`no_std`) compose with autodiff only. See
[Supported Backends](#supported-backends) below for the full matrix.
Every project here is open-source and actively developed. Want to help build the Rust AI ecosystem?
The [good first issues](https://github.com/tracel-ai/burn/contribute) are a great place to start,
and the [Contributing](#contributing) guide will get you set up.
<details>
<summary>
<b>Community crates 🌱</b>
</summary>
<br />
These crates are not maintained by Tracel, but they are part of the same Rust AI story. Anything
that helps you load data, build environments, or ship models belongs here. Built something that
fits? Open a PR to add it!
| Category | Crate | Description |
| -------------------------- | --------------------------------------------------------------- | ----------------------------------------------------------------- |
| Data & loading | [polars](https://github.com/pola-rs/polars) | Fast DataFrames for tabular data |
| | [arrow-rs](https://github.com/apache/arrow-rs) | Apache Arrow columnar memory format |
| | [image](https://github.com/image-rs/image) | Image decoding, encoding and processing |
| | [hf-hub](https://github.com/huggingface/hf-hub) | Download models and datasets from the Hugging Face Hub |
| Tokenization & NLP | [tokenizers](https://github.com/huggingface/tokenizers) | Fast, production-ready tokenizers |
| | [rust-bert](https://github.com/guillaume-be/rust-bert) | Ready-to-use NLP pipelines and transformer models |
| Numerical & linear algebra | [ndarray](https://github.com/rust-ndarray/ndarray) | N-dimensional arrays |
| | [nalgebra](https://github.com/dimforge/nalgebra) | Linear algebra |
| Classical ML | [linfa](https://github.com/rust-ml/linfa) | Classical ML toolkit, in the spirit of scikit-learn |
| | [smartcore](https://github.com/smartcorelib/smartcore) | Classical ML algorithms, no BLAS/LAPACK required |
| Inference & runtimes | [candle](https://github.com/huggingface/candle) | Minimalist ML framework with a focus on LLM inference |
| | [mistral.rs](https://github.com/EricLBuehler/mistral.rs) | Fast, multimodal LLM inference engine |
| | [ort](https://github.com/pykeio/ort) | ONNX Runtime bindings for hardware-accelerated inference |
| | [tract](https://github.com/sonos/tract) | Pure-Rust inference for ONNX and NNEF models |
| | [wonnx](https://github.com/webonnx/wonnx) | 100% Rust, WebGPU-accelerated ONNX runtime for native and the web |
| LLM apps & RAG | [rig](https://github.com/0xPlaygrounds/rig) | Build modular LLM applications and agents |
| | [langchain-rust](https://github.com/Abraxas-365/langchain-rust) | LangChain-style chain orchestration |
| Embeddings & vector search | [fastembed](https://github.com/Anush008/fastembed-rs) | Generate text embeddings and rerank locally |
| | [qdrant](https://github.com/qdrant/qdrant) | Vector search engine, written in Rust |
| | [lancedb](https://github.com/lancedb/lancedb) | Embedded, developer-friendly vector database |
| Computer vision | [kornia-rs](https://github.com/kornia/kornia-rs) | Low-level 3D computer vision library |
| Simulation & environments | [rapier](https://github.com/dimforge/rapier) | Physics engine for robotics and RL environments |
| Visualization | [rerun](https://github.com/rerun-io/rerun) | Multimodal data and CV/robotics visualization |
| | [plotters](https://github.com/plotters-rs/plotters) | Plotting and charting |
</details>
## Backend
<div align="left">
<img align="right" src="https://raw.githubusercontent.com/tracel-ai/burn/main/assets/backend-chip.png" height="96px"/>
Burn strives to be as fast as possible on as many hardwares as possible, with robust
implementations. We believe this flexibility is crucial for modern needs where you may train your
models in the cloud, then deploy on customer hardwares, which vary from user to user.
</div>
### Supported Backends
Most backends support all operating systems, so we don't mention them in the tables below.
**GPU Backends:**
| | CUDA | ROCm | Metal | Vulkan | WebGPU | LibTorch |
| ------- | ---- | ---- | ----- | ------ | ------ | -------- |
| Nvidia | ☑️ | - | - | ☑️ | ☑️ | ☑️ |
| AMD | - | ☑️ | - | ☑️ | ☑️ | ☑️ |
| Apple | - | - | ☑️ | - | ☑️ | ☑️ |
| Intel | - | - | - | ☑️ | ☑️ | - |
| Qualcom | - | - | - | ☑️ | ☑️ | - |
| Wasm | - | - | - | - | ☑️ | - |
**CPU Backends:**
| | Cpu (CubeCL) | Flex | LibTorch |
| ------ | ------------ | ---- | -------- |
| X86 | ☑️ | ☑️ | ☑️ |
| Arm | ☑️ | ☑️ | ☑️ |
| Wasm | - | ☑️ | - |
| no-std | - | ☑️ | - |
<br />
Compared to other frameworks, Burn has a very different approach to supporting many backends. By
design, most code is generic over the Backend trait, which allows us to build Burn with swappable
backends. This makes composing backend possible, augmenting them with additional functionalities
such as autodifferentiation and automatic kernel fusion.
<details>
<summary>
Autodiff: Backend decorator that brings backpropagation to any backend 🔄
</summary>
<br />
Contrary to the aforementioned backends, Autodiff is actually a backend _decorator_. This means that
it cannot exist by itself; it must encapsulate another backend.
The simple act of wrapping a base backend with Autodiff transparently equips it with
autodifferentiation support, making it possible to call backward on your model.
```rust
use burn::backend::{Autodiff, Wgpu};
use burn::tensor::{Distribution, Tensor};
fn main() {
type Backend = Autodiff<Wgpu>;
let device = Default::default();
let x: Tensor<Backend, 2> = Tensor::random([32, 32], Distribution::Default, &device);
let y: Tensor<Backend, 2> = Tensor::random([32, 32], Distribution::Default, &device).require_grad();
let tmp = x.clone() + y.clone();
let tmp = tmp.matmul(x);
let tmp = tmp.exp();
let grads = tmp.backward();
let y_grad = y.grad(&grads).unwrap();
println!("{y_grad}");
}
```
Of note, it is impossible to make the mistake of calling backward on a model that runs on a backend
that does not support autodiff (for inference), as this method is only offered by an Autodiff
backend.
See the [Autodiff Backend README](./crates/burn-autodiff/README.md) for more details.
</details>
<details>
<summary>
Fusion: Backend decorator that brings kernel fusion to all first-party backends
</summary>
<br />
This backend decorator enhances a backend with kernel fusion, provided that the inner backend
supports it. Note that you can compose this backend with other backend decorators such as Autodiff.
All first-party accelerated backends (like WGPU and CUDA) use Fusion by default (`burn/fusion`
feature flag), so you typically don't need to apply it manually.
```rust
#[cfg(not(feature = "fusion"))]
pub type Cuda<F = f32, I = i32> = CubeBackend<CudaRuntime, F, I, u8>;
#[cfg(feature = "fusion")]
pub type Cuda<F = f32, I = i32> = burn_fusion::Fusion<CubeBackend<CudaRuntime, F, I, u8>>;
```
Of note, we plan to implement automatic gradient checkpointing based on compute bound and memory
bound operations, which will work gracefully with the fusion backend to make your code run even
faster during training, see [this issue](https://github.com/tracel-ai/burn/issues/936).
See the [Fusion Backend README](./crates/burn-fusion/README.md) for more details.
</details>
<details>
<summary>
Remote (Beta): Backend decorator for remote backend execution, useful for distributed computations
</summary>
<br />
That backend has two parts, one client and one server. The client sends tensor operations over the
network to a remote compute backend. You can use any first-party backend as server in a single line
of code:
```rust
fn main_server() {
// Start a server on port 3000.
burn::server::start::<burn::backend::Cuda>(Default::default(), 3000);
}
fn main_client() {
// Create a client that communicate with the server on port 3000.
use burn::backend::{Autodiff, RemoteBackend};
type Backend = Autodiff<RemoteDevice>;
let device = RemoteDevice::new("ws://localhost:3000");
let tensor_gpu =
Tensor::<Backend, 2>::random([3, 3], Distribution::Default, &device);
}
```
</details>
<br />
## Training & Inference
<div align="left">
<img align="right" src="https://raw.githubusercontent.com/tracel-ai/burn/main/assets/ember-wall.png" height="96px"/>
The whole deep learning workflow is made easy with Burn, as you can monitor your training progress
with an ergonomic dashboard, and run inference everywhere from embedded devices to large GPU
clusters.
Burn was built from the ground up with training and inference in mind. It's also worth noting how
Burn, in comparison to frameworks like PyTorch, simplifies the transition from training to
deployment, eliminating the need for code changes.
</div>
<div align="center">
<br />
<a href="https://www.youtube.com/watch?v=N9RM5CQbNQc" target="_blank">
<img src="https://raw.githubusercontent.com/tracel-ai/burn/main/assets/burn-train-tui.png" alt="Burn Train TUI" width="75%">
</a>
</div>
<br />
**Click on the following sections to expand 👇**
<details>
<summary>
Training Dashboard 📈
</summary>
<br />
As you can see in the previous video (click on the picture!), a new terminal UI dashboard based on
the [Ratatui](https://github.com/ratatui-org/ratatui) crate allows users to follow their training
with ease without having to connect to any external application.
You can visualize your training and validation metrics updating in real-time and analyze the
lifelong progression or recent history of any registered metrics using only the arrow keys. Break
from the training loop without crashing, allowing potential checkpoints to be fully written or
important pieces of code to complete without interruption 🛡
</details>
<details>
<summary>
ONNX Support 🐫
</summary>
<br />
Burn supports importing ONNX (Open Neural Network Exchange) models through the
[burn-onnx](https://github.com/tracel-ai/burn-onnx) crate, allowing you to easily port models from
TensorFlow or PyTorch to Burn. The ONNX model is converted into Rust code that uses Burn's native
APIs, enabling the imported model to run on any Burn backend (CPU, GPU, WebAssembly) and benefit
from all of Burn's optimizations like automatic kernel fusion.
Our ONNX support is further described in
[this section of the Burn Book 🔥](https://burn.dev/books/burn/onnx-import.html).
> **Note**: This crate is in active development and currently supports a
> [limited set of ONNX operators](https://github.com/tracel-ai/burn-onnx/blob/main/SUPPORTED-ONNX-OPS.md).
</details>
<details>
<summary>
Importing PyTorch or Safetensors Models 🚚
</summary>
<br />
You can load weights from PyTorch or Safetensors formats directly into your Burn-defined models.
This makes it easy to reuse existing models while benefiting from Burn's performance and deployment
features.
Learn more in the [Saving & Loading Models](https://burn.dev/books/burn/saving-and-loading.html)
section of the Burn Book.
</details>
<details>
<summary>
Inference in the Browser 🌐
</summary>
<br />
Several of our backends can run in WebAssembly environments: Flex for CPU execution, and WGPU for
GPU acceleration via WebGPU. This means that you can run inference directly within a browser. We
provide several examples of this:
- [MNIST](./examples/mnist-inference-web) where you can draw digits and a small convnet tries to
find which one it is! 2️⃣ 7️⃣ 😰
- [Image Classification](https://github.com/tracel-ai/burn-onnx/tree/main/examples/image-classification-web)
where you can upload images and classify them! 🌄
</details>
<details>
<summary>
Embedded: <i>no_std</i> support ⚙️
</summary>
<br />
Burn's core components support [no_std](https://docs.rust-embedded.org/book/intro/no-std.html). This
means it can run in bare metal environment such as embedded devices without an operating system.
> As of now, only the Flex backend can be used in a _no_std_ environment.
</details>
<br />
### Benchmarks
To evaluate performance across different backends and track improvements over time, we provide a
dedicated benchmarking suite.
Run and compare benchmarks using [burn-bench](https://github.com/tracel-ai/burn-bench).
> ⚠️ **Warning** When using one of the `wgpu` backends, you may encounter compilation errors related
> to recursive type evaluation. This is due to complex type nesting within the `wgpu` dependency
> chain. To resolve this issue, add the following line at the top of your `main.rs` or `lib.rs`
> file:
>
> ```rust
> #![recursion_limit = "256"]
> ```
>
> The default recursion limit (128) is often just below the required depth (typically 130-150) due
> to deeply nested associated types and trait bounds.
## Getting Started
<div align="left">
<img align="right" src="https://raw.githubusercontent.com/tracel-ai/burn/main/assets/ember-walking.png" height="96px"/>
Just heard of Burn? You are at the right place! Just continue reading this section and we hope you
can get on board really quickly.
</div>
<details>
<summary>
The Burn Book 🔥
</summary>
<br />
To begin working effectively with Burn, it is crucial to understand its key components and
philosophy. This is why we highly recommend new users to read the first sections of
[The Burn Book 🔥](https://burn.dev/books/burn/). It provides detailed examples and explanations
covering every facet of the framework, including building blocks like tensors, modules, and
optimizers, all the way to advanced usage, like coding your own GPU kernels.
> The project is constantly evolving, and we try as much as possible to keep the book up to date
> with new additions. However, we might miss some details sometimes, so if you see something weird,
> let us know! We also gladly accept Pull Requests 😄
</details>
<details>
<summary>
Examples 🙏
</summary>
<br />
Let's start with a code snippet that shows how intuitive the framework is to use! In the following,
we declare a neural network module with some parameters along with its forward pass.
```rust
use burn::nn;
use burn::module::Module;
use burn::tensor::backend::Backend;
#[derive(Module, Debug)]
pub struct PositionWiseFeedForward<B: Backend> {
linear_inner: nn::Linear<B>,
linear_outer: nn::Linear<B>,
dropout: nn::Dropout,
gelu: nn::Gelu,
}
impl<B: Backend> PositionWiseFeedForward<B> {
pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
let x = self.linear_inner.forward(input);
let x = self.gelu.forward(x);
let x = self.dropout.forward(x);
self.linear_outer.forward(x)
}
}
```
We have a somewhat large amount of [examples](./examples) in the repository that shows how to use
the framework in different scenarios.
Following [the book](https://burn.dev/books/burn/):
- [Basic Workflow](./examples/guide) : Creates a custom CNN `Module` to train on the MNIST dataset
and use for inference.
- [Custom Training Loop](./examples/custom-training-loop) : Implements a basic training loop instead
of using the `Learner`.
- [Custom WGPU Kernel](./examples/custom-wgpu-kernel) : Learn how to create your own custom
operation with the WGPU backend.
Additional examples:
- [Custom CSV Dataset](./examples/custom-csv-dataset) : Implements a dataset to parse CSV data for a
regression task.
- [Regression](./examples/simple-regression) : Trains a simple MLP on the California Housing dataset
to predict the median house value for a district.
- [Custom Image Dataset](./examples/custom-image-dataset) : Trains a simple CNN on custom image
dataset following a simple folder structure.
- [Custom Renderer](./examples/custom-renderer) : Implements a custom renderer to display the
[`Learner`](./building-blocks/learner.md) progress.
- [Image Classification Web](./examples/image-classification-web) : Image classification web browser
demo using Burn, WGPU and WebAssembly.
- [MNIST Inference on Web](./examples/mnist-inference-web) : An interactive MNIST inference demo in
the browser. The demo is available [online](https://burn.dev/demo/).
- [MNIST Training](./examples/mnist) : Demonstrates how to train a custom `Module` (MLP) with the
`Learner` configured to log metrics and keep training checkpoints.
- [PyTorch Import Inference](./examples/import-model-weights) : Imports a PyTorch model pre-trained
on MNIST to perform inference on a sample image with Burn.
- [Text Classification](./examples/text-classification) : Trains a text classification transformer
model on the AG News or DbPedia dataset. The trained model can then be used to classify a text
sample.
- [Text Generation](./examples/text-generation) : Trains a text generation transformer model on the
DbPedia dataset.
- [Wasserstein GAN MNIST](./examples/wgan) : Trains a WGAN model to generate new handwritten digits
based on MNIST.
For more practical insights, you can clone the repository and run any of them directly on your
computer!
</details>
<details>
<summary>
Pre-trained Models 🤖
</summary>
<br />
We keep an updated and curated list of models and examples built with Burn, see the
[tracel-ai/models repository](https://github.com/tracel-ai/models) for more details.
Don't see the model you want? Don't hesitate to open an issue, and we may prioritize it. Built a
model using Burn and want to share it? You can also open a Pull Request and add your model under the
community section!
</details>
<details>
<summary>
Why use Rust for AI? 🦀
</summary>
<br />
Deep Learning is a special form of software where you need very high level abstractions as well as
extremely fast execution time. Rust is the perfect candidate for that use case since it provides
zero-cost abstractions to easily create neural network modules, and fine-grained control over memory
to optimize every detail. To this day, the mainstream solution has been to offer APIs in Python but
rely on bindings to low-level languages such as C/C++. This reduces portability, increases
complexity and creates friction between researchers and engineers. Rust's approach to abstractions
is versatile enough to tackle this two-language dichotomy, and Cargo makes it easy to build, test
and deploy from any environment, which is usually a pain in Python.
Rust's AI ecosystem is young, but it is real and growing quickly. Foundational pieces are already
here: Burn and [CubeCL](https://github.com/tracel-ai/cubecl) for training and compute,
[candle](https://github.com/huggingface/candle) for inference, Hugging Face's `tokenizers` and
`safetensors`, and `polars` and `ndarray` for data. Betting on Rust today means betting on a stack
that is growing, and one where contributors still shape the direction. The pieces that don't exist
yet are opportunities rather than dead-ends (see [Contributing](#contributing)).
Rust is also what makes one-stack-everywhere possible: a single self-contained binary with no Python
runtime to ship, running from servers down to `no_std` embedded targets.
</details>
<br />
> **Deprecation Note**<br />Since `0.14.0`, the internal structure for tensor data has changed. The
> previous `Data` struct was deprecated and officially removed since `0.17.0` in favor of the new
> `TensorData` struct, which allows for more flexibility by storing the underlying data as bytes and
> keeping the data type as a field. If you are using `Data` in your code, make sure to switch to
> `TensorData`.
<!-- >
> In the event that you are trying to load a model record saved in a previous version, make sure to
> enable the `record-backward-compat` feature using a previous version of burn (<=0.16.0). Otherwise,
> the record won't be deserialized correctly and you will get an error message (which will also point
> you to the backward compatible feature flag). The backward compatibility was maintained for
> deserialization (loading), so as soon as you have saved the record again it will be saved according
> to the new structure and you will be able to upgrade to this version. Please note that binary formats
> are not backward compatible. Thus, you will need to load your record in a previous version and save it
> to another of the self-describing record formats before using a compatible version (as described) with the
> `record-backward-compat` feature flag. -->
<details id="deprecation">
<summary>
Loading Model Records From Previous Versions ⚠️
</summary>
<br />
In the event that you are trying to load a model record saved in a version older than `0.14.0`, make
sure to use a compatible version (`0.14`, `0.15` or `0.16`) with the `record-backward-compat`
feature flag.
```
features = [..., "record-backward-compat"]
```
Otherwise, the record won't be deserialized correctly and you will get an error message. This error
will also point you to the backward compatible feature flag.
The backward compatibility was maintained for deserialization when loading records. Therefore, as
soon as you have saved the record again it will be saved according to the new structure and you can
upgrade back to the current version
Please note that binary formats are not backward compatible. Thus, you will need to load your record
in a previous version and save it in any of the other self-describing record format (e.g., using the
`NamedMpkFileRecorder`) before using a compatible version (as described) with the
`record-backward-compat` feature flag.
</details>
## Community
<div align="left">
<img align="right" src="https://raw.githubusercontent.com/tracel-ai/burn/main/assets/ember-community.png" height="96px"/>
If you are excited about the project, don't hesitate to join our
[Discord](https://discord.gg/uPEBbYYDB6)! We try to be as welcoming as possible to everybody from
any background. You can ask your questions and share what you built with the community!
</div>
<br/>
### Contributing
Before contributing, please read the [Contributing Guidelines](./CONTRIBUTING.md) and our
[Code of Conduct](./CODE-OF-CONDUCT.md). The [Contributor Book](https://burn.dev/contributor-book/)
covers architecture, environment setup, and guides for common tasks.
## Status
Burn is currently in active development, and there will be breaking changes. While any resulting
issues are likely to be easy to fix, there are no guarantees at this stage.
## License
Burn is distributed under the terms of both the MIT license and the Apache License (Version 2.0).
See [LICENSE-APACHE](./LICENSE-APACHE) and [LICENSE-MIT](./LICENSE-MIT) for details. Opening a pull
request is assumed to signal agreement with these licensing terms.
</div>
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`tracel-ai/burn`
- 原始仓库：https://github.com/tracel-ai/burn
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
_typos.toml
+24
View File
@@ -0,0 +1,24 @@
[default]
extend-ignore-identifiers-re = ["ratatui", "Ratatui", "NdArray*", "ND", "log_lik"]
[default.extend-identifiers]
UE4M3 = "UE4M3"
UE8M0 = "UE8M0"
ue8m0 = "ue8m0"
[files]
extend-exclude = [
"*.onnx",
"*.proto",
"assets/ModuleSerialization.xml",
]
[default.extend-words]
# Don't correct "nd" (n-dimensional, as in scatter_nd/gather_nd)
nd = "nd"
# Don't correct "arange" which is intentional
arange = "arange"
# Don't correct "convnet" (convolutional network)
convnet = "convnet"
# Don't correct "Nd" in convNd / conv_nd (N-dimensional)
Nd = "Nd"
assets/backend-chip.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 67 KiB

assets/burn-train-tui.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 985 KiB

assets/ember-blazingly-fast.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 165 KiB

assets/ember-community.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 122 KiB

assets/ember-walking.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 133 KiB

assets/ember-wall.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 173 KiB

assets/logo-burn-full.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 139 KiB

assets/logo-burn-neutral.webp
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 28 KiB

assets/logo-burn-small.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 62 KiB

benchmarks.toml
+82
View File
@@ -0,0 +1,82 @@
[environment]
gcp_gpu_attached = true
gcp_image_family = "tracel-ci-ubuntu-2404-amd64-nvidia"
# https://cloud.google.com/compute/docs/accelerator-optimized-machines
# put the faster machine on first place for possibly faster 'Benchmarks Started' feedback in PRs
gcp_machine_types = [
"a2-highgpu-1g", # 1 A100 40GB (listed as a2 standard)
"g2-standard-4", # 1 L4 24GB
]
# define the available zones for each machine type
# be sure to check what machine types are available in each region
# https://cloud.google.com/compute/docs/gpus/gpu-regions-zones#view-using-table
gcp_zones = [
# a2-highgpu-1g
[
"asia-northeast1-a",
"asia-northeast1-c",
"asia-northeast3-b",
"asia-southeast1-b",
"asia-southeast1-c",
"europe-west4-a",
"europe-west4-b",
"us-central1-a",
"us-central1-b",
"us-central1-c",
"us-central1-f",
"us-east1-b",
"us-west1-b",
"us-west3-b",
"us-west4-b"
],
# g2-standard-4
[
"northamerica-northeast2-a",
"northamerica-northeast2-b",
"us-central1-a",
"us-central1-b",
"us-central1-c",
"us-east1-b",
"us-east1-c",
"us-east1-d",
"us-east4-a",
"us-east4-c",
"us-west1-a",
"us-west1-b",
"us-west1-c",
"us-west4-a",
"us-west4-c"
],
]
repo_full = "tracel-ai/burn"
rust_toolchain = "stable"
rust_version = "stable"
[burn-bench]
github_organization = "tracel-ai"
github_repository = "burn-bench"
github_branch = "main"
github_workflow = "benchmarks.yml"
# vulkan autotune seems to take ages, disabling it for now
# backends = ["cuda-fusion", "vulkan-fusion", "wgpu-fusion"]
backends = ["cuda-fusion", "cuda"]
benches = ["autodiff",
"binary",
"bool_select",
"conv-transpose2d",
"conv-transpose3d",
"conv2d",
"conv3d",
"custom-gelu",
"data",
"load-record",
"matmul-fused",
"matmul",
"max-pool2d",
"random",
"reduce",
"softmax",
"transformer-encoder",
"unary"
]
dtypes = ["f16"]
burn-book/.gitignore
+18
View File
@@ -0,0 +1,18 @@
target
# MacOS temp file
.DS_Store
book-test
guide/book
.vscode
tests/burn-book/book/
book/
# Ignore Jetbrains specific files.
.idea/
# Ignore Vim temporary and swap files.
*.sw?
*~
burn-book/.prettierrc.json
+4
View File
@@ -0,0 +1,4 @@
{
"printWidth": 100,
"proseWrap": "always"
}
burn-book/LICENSE-APACHE
Symlink
+1
View File
@@ -0,0 +1 @@
../LICENSE-APACHE
burn-book/LICENSE-MIT
Symlink
+1
View File
@@ -0,0 +1 @@
../LICENSE-MIT
burn-book/book.toml
+16
View File
@@ -0,0 +1,16 @@
[book]
authors = [
"Wouter Doppenberg",
"Nathaniel Simard",
"Louis Fortier-Dubois",
"Dilshod Tadjibaev",
"Guillaume Lagrange",
"Sylvain Benner",
"Bjorn Beishline"
]
language = "en"
src = "src"
title = "The Burn Book 🔥"
[output.html]
mathjax-support = true
burn-book/src/SUMMARY.md
+38
View File
@@ -0,0 +1,38 @@
- [Overview](./overview.md)
- [Why Burn?](./motivation.md)
- [Getting started](./getting-started.md)
- [Examples](./examples.md)
- [Basic Workflow: From Training to Inference](./basic-workflow/README.md)
- [Model](./basic-workflow/model.md)
- [Data](./basic-workflow/data.md)
- [Training](./basic-workflow/training.md)
- [Backend](./basic-workflow/backend.md)
- [Inference](./basic-workflow/inference.md)
- [Building Blocks](./building-blocks/README.md)
- [Backend](./building-blocks/backend.md)
- [Tensor](./building-blocks/tensor.md)
- [Autodiff](./building-blocks/autodiff.md)
- [Module](./building-blocks/module.md)
- [Learner](./building-blocks/learner.md)
- [Metric](./building-blocks/metric.md)
- [Config](./building-blocks/config.md)
- [Record](./building-blocks/record.md)
- [Dataset](./building-blocks/dataset.md)
- [Performance](./performance/README.md)
- [Good practices](./performance/good-practices/README.md)
- [Asynchronous Execution](./performance/good-practices/asynchronous-execution.md)
- [Kernel Fusion](./performance/good-practices/kernel-fusion.md)
- [Kernel Selection](./performance/good-practices/kernel-selection.md)
- [Quantization](./performance/quantization.md)
- [Distributed Computing](./performance/distributed-computing.md)
- [Custom Training Loop](./custom-training-loop.md)
- [Saving & Loading Models](./saving-and-loading.md)
- [ONNX Import](./onnx-import.md)
- [Models & Pre-Trained Weights](./models-and-pretrained-weights.md)
- [Advanced](./advanced/README.md)
- [Backend Extension](./advanced/backend-extension/README.md)
- [Custom `CubeCL` Kernel](./advanced/backend-extension/custom-cubecl-kernel.md)
- [Custom WGPU Kernel](./advanced/backend-extension/custom-wgpu-kernel.md)
- [Custom Optimizer]()
- [WebAssembly](./advanced/web-assembly.md)
- [No-Std](./advanced/no-std.md)
burn-book/src/advanced/README.md
+12
View File
@@ -0,0 +1,12 @@
# Advanced
In this section, we will go into advanced topics that extend beyond basic usage. Given Burn's
exceptional flexibility, a lot of advanced use cases become possible.
Before going through this section, we strongly recommend exploring the
[basic workflow](../basic-workflow/) section and the
[building blocks](../building-blocks/) section. Establishing a solid understanding of how
the framework operates is crucial to comprehending the advanced concepts presented here. While you
have the freedom to explore the advanced sections in any order you prefer, it's important to note
that this section is not intended to be linear, contrary to preceding sections. Instead, it serves
as a repository of use cases that you can refer to for guidance as needed.
burn-book/src/advanced/backend-extension/README.md
+82
View File
@@ -0,0 +1,82 @@
# Backend Extension
Burn aims to be the most flexible deep learning framework. While it's crucial to maintain
compatibility with a wide variety of backends, Burn provides the ability to extend the functionality
of a backend implementation to suit your modeling requirements. This versatility is advantageous in
numerous ways, such as supporting custom operations like flash attention or manually fusing
operations for enhanced performance.
In this section, we will go into the process of extending a backend, providing multiple examples.
But before we proceed, let's establish the fundamental principles that will empower you to craft
your own backend extensions.
As you can observe, most types in Burn are generic over the Backend trait. This might give the
impression that Burn operates at a high level over the backend layer. However, making the trait
explicit instead of being chosen via a compilation flag was a thoughtful design decision. This
explicitness does not imply that all backends must be identical; rather, it offers a great deal of
flexibility when composing backends. The autodifferentiation backend trait (see
[autodiff section](../../building-blocks/autodiff.md)) is an example of how the backend trait has
been extended to enable gradient computation with backpropagation. Furthermore, this design allows
you to create your own backend extension. To achieve this, you need to design your own backend trait
specifying which functions should be supported.
```rust, ignore
pub trait Backend: burn::tensor::backend::Backend {
fn my_new_function(tensor: B::FloatTensorPrimitive) -> B::FloatTensorPrimitive {
// You can define a basic implementation reusing the Burn Backend API.
// This can be useful since all backends will now automatically support
// your model. But performance can be improved for this new
// operation by implementing this block in specific backends.
}
}
```
You can then implement your new custom backend trait for any backend that you want to support:
```rust, ignore
impl Backend for burn_tch::LibTorch {
fn my_new_function(tensor: TchTensor) -> TchTensor {
// My Tch implementation
}
}
impl Backend for burn_flex::Flex {
// No specific implementation, but the backend can still be used.
}
```
You can support the backward pass using the same pattern.
```rust, ignore
impl<B: Backend> Backend for burn_autodiff::Autodiff<B> {
// No specific implementation; autodiff will work with the default
// implementation. Useful if you still want to train your model, but
// observe performance gains mostly during inference.
}
impl<B: Backend> Backend for burn_autodiff::Autodiff<B> {
fn my_new_function(tensor: AutodiffTensor) -> AutodiffTensor {
// My own backward implementation, generic over my custom Backend trait.
//
// You can add a new method `my_new_function_backward` to your custom backend
// trait if you want to invoke a custom kernel during the backward pass.
}
}
impl Backend for burn_autodiff::Autodiff<burn_tch::LibTorch> {
fn my_new_function(tensor: AutodiffTensor) -> AutodiffTensor {
// My own backward implementation, generic over a backend implementation.
//
// This is another way to call a custom kernel for the backward pass that
// doesn't require the addition of a new `backward` function in the custom backend.
// This is useful if you don't want all backends to support training, reducing
// the need for extra code when you know your model will only be trained on one
// specific backend.
}
}
```
The specifics of each implementation will be covered by the examples provided in this section. The
`cubecl` compiler frontend is the recommended method of implementing custom kernels, since it
supports multiple backends, including `wgpu` and `CUDA`, and is the way first-party `burn` kernels
are written.
burn-book/src/advanced/backend-extension/custom-cubecl-kernel.md
+383
View File
@@ -0,0 +1,383 @@
# Custom CubeCL Kernel
In this section, you will learn how to create your own custom operation by writing your own kernel
with the cubecl compiler frontend. We will take the example of a common workflow in the deep
learning field, where we create a kernel to fuse multiple operations together. Note that `burn` does
this automatically, but a manual implementation might be more efficient in some cases. We will fuse
a matmul kernel followed by an addition and the ReLU activation function, which is commonly found in
various models. All the code can be found under the
[examples directory](https://github.com/tracel-ai/burn/tree/main/examples/custom-cubecl-kernel).
> Note: CubeCL is in active development, so this section may be outdated.
## Custom Backend Trait
First, we need to determine the type signature of our newly created operation by defining our custom
backend traits. As we will use the associated type `TensorPrimitive` of the `Backend` trait, which
encapsulates the underlying tensor implementation of the backend, we will use a type alias to avoid
the ugly disambiguation with associated types.
```rust, ignore
/// We create our own Backend trait that extends the Burn backend trait.
pub trait Backend: burn::tensor::backend::Backend {
fn fused_matmul_add_relu(
lhs: FloatTensor<Self>,
rhs: FloatTensor<Self>,
bias: FloatTensor<Self>,
) -> FloatTensor<Self>;
}
/// We create our own AutodiffBackend trait that extends the Burn autodiff backend trait.
pub trait AutodiffBackend: Backend + burn::tensor::backend::AutodiffBackend {}
```
In our project, we can use these traits instead of the
`burn::tensor::backend::{Backend, AutodiffBackend}` traits provided by Burn. Burn's user APIs
typically make use of the `Tensor` struct rather than dealing directly with primitive tensor types.
Therefore, we can encapsulate our newly defined backend traits with functions that expose new
operations while maintaining a consistent API.
```rust, ignore
/// We define our custom implementation using the added function on our custom backend.
pub fn matmul_add_relu_custom<B: Backend>(
lhs: Tensor<B, 3>,
rhs: Tensor<B, 3>,
bias: Tensor<B, 3>,
) -> Tensor<B, 3> {
let output = B::fused_matmul_add_relu(
lhs.into_primitive().tensor(),
rhs.into_primitive().tensor(),
bias.into_primitive().tensor(),
);
Tensor::from_primitive(TensorPrimitive::Float(output))
}
/// We define a reference implementation using basic tensor operations.
pub fn matmul_add_relu_reference<B: Backend>(
lhs: Tensor<B, 3>,
rhs: Tensor<B, 3>,
bias: Tensor<B, 3>,
) -> Tensor<B, 3> {
let x = lhs.matmul(rhs) + bias;
activation::relu(x)
}
```
Note that we also provide a reference implementation for testing purposes, which allows us to easily
validate our new implementation. While not mandatory, having a reference implementation can be
valuable, especially in projects where creating a reference implementation solely using basic tensor
operations is feasible.
## Forward Kernel
Now, let's proceed to write the fused kernel using the `cubecl` compiler frontend. To keep things
simple, we'll create a straightforward matmul kernel without employing any intricate techniques. We
won't delve into the details of the `cube` macro, but if you're interested to learn more, please see
[`cubecl` Book](https://github.com/tracel-ai/cubecl/tree/f5b63076a01a5c03ea9ed20799d3eeaf776b45da/cubecl-book).
The actual matmul, add and relu computations are found at the end, after an extensive prelude that
serves to correctly map each compute unit to the data it is responsible for, with support for
batches.
```rust, ignore
use cubecl::{cube, prelude::*};
#[cube(launch)]
pub fn fused_matmul_add_relu_kernel<F: Float>(
lhs: &Tensor<F>,
rhs: &Tensor<F>,
bias: &Tensor<F>,
output: &mut Tensor<F>,
) {
let row = ABSOLUTE_POS_X;
let col = ABSOLUTE_POS_Y;
let batch = ABSOLUTE_POS_Z;
let n_rows = output.shape(output.rank() - 2);
let n_cols = output.shape(output.rank() - 1);
let dim_k = rhs.shape(rhs.rank() - 1);
if row >= n_rows || col >= n_cols {
return;
}
let offset_output = batch * n_rows * n_cols;
let mut offset_lhs = 0;
let mut offset_rhs = 0;
let batch_dims = output.rank() - 2;
for dim in 0..batch_dims {
offset_lhs += offset_output / output.stride(dim) % lhs.shape(dim) * lhs.stride(dim);
offset_rhs += offset_output / output.stride(dim) % rhs.shape(dim) * rhs.stride(dim);
}
let mut sum = F::new(0.0);
for k in 0..dim_k {
let lhs_index = row * dim_k + k;
let rhs_index = k * n_cols + col;
sum += lhs[offset_lhs + lhs_index] * rhs[offset_rhs + rhs_index];
}
let out_index = row * n_cols + col;
let index = offset_output + out_index;
output[index] = F::max(sum + bias[index], F::new(0.0));
}
```
Now, let's move on to the next step, which involves implementing the remaining code to launch the
kernel. We'll go into implementing our custom backend trait for the generic JIT backend. This
automatically implements the trait for `burn-cuda`, `burn-wgpu` as well as fusion.
```rust, ignore
/// Implement our custom backend trait for the generic `CubeBackend`.
impl<R: CubeRuntime> Backend for CubeBackend<R>
{
fn fused_matmul_add_relu(
lhs: FloatTensor<Self>,
rhs: FloatTensor<Self>,
bias: FloatTensor<Self>,
) -> FloatTensor<Self> {
// Define cube dim, hardcoded for simplicity.
let cube_dim = CubeDim { x: 16, y: 16, z: 1 };
lhs.assert_is_on_same_device(&rhs);
lhs.assert_is_on_same_device(&bias);
// For simplicity, make sure each tensor is continuous.
let lhs = into_contiguous(lhs);
let rhs = into_contiguous(rhs);
let bias = into_contiguous(bias);
// Get the matmul relevant shapes.
let ndims = lhs.shape.num_dims();
let num_rows = lhs.shape[ndims - 2];
let num_cols = rhs.shape[ndims - 1];
// Compute shape of output, while tracking number of batches.
let mut num_batches = 1;
let mut shape_out = vec![0; ndims];
for i in shape_out.clone().into_iter().take(ndims - 2) {
shape_out[i] = usize::max(lhs.shape[i], rhs.shape[i]);
num_batches *= shape_out[i];
}
shape_out[ndims - 2] = num_rows;
shape_out[ndims - 1] = num_cols;
let shape_out = Shape::from(shape_out);
// Create a buffer for the output tensor.
let buffer = lhs
.client
.empty(shape_out.num_elements() * core::mem::size_of::<F>());
// Create the output tensor primitive.
let output = CubeTensor::new_contiguous(
lhs.client.clone(),
lhs.device.clone(),
shape_out,
buffer,
F::dtype(),
);
// Declare the wgsl workgroup with the number of cubes in x, y and z.
let cubes_needed_in_x = f32::ceil(num_rows as f32 / cube_dim.x as f32) as u32;
let cubes_needed_in_y = f32::ceil(num_cols as f32 / cube_dim.y as f32) as u32;
let cube_count =
CubeCount::Static(cubes_needed_in_x, cubes_needed_in_y, num_batches as u32);
// Execute lazily the kernel with the launch information and the given buffers. For
// simplicity, no vectorization is performed
fused_matmul_add_relu_kernel::launch::<F, R>(
&lhs.client,
cube_count,
cube_dim,
lhs.into_tensor_arg(),
rhs.into_tensor_arg(),
bias.into_tensor_arg(),
output.clone().into_tensor_arg(),
);
// Return the output tensor.
output
}
}
```
In the preceding code block, we demonstrated how to launch the kernel that modifies the correct
buffer. It's important to note that Rust's mutability safety doesn't apply here; the context has the
capability to execute any mutable operation on any buffer. While this isn't a problem in the
previous scenario where we only modify the newly created output buffer, it is wise to keep this in
mind.
## Backward
Now that the custom backend trait is implemented for the JIT backend, you can use it to invoke the
`matmul_add_relu_custom` function. However, calculating gradients is not yet possible at this stage.
If your use case does not extend beyond inference, there is no need to implement any of the
following code.
For the backward pass, we will leverage the backend implementation from `burn-autodiff`, which is
actually generic over the backend. Instead of crafting our own `cubecl` kernel for the backward
pass, we will use our fused kernel only for the forward pass, and compute the gradient using basic
operations.
```rust, ignore
// Implement our custom backend trait for any backend that also implements our custom backend trait.
impl<B: Backend, C: CheckpointStrategy> Backend for Autodiff<B, C> {
fn fused_matmul_add_relu(
lhs: FloatTensor<Self>,
rhs: FloatTensor<Self>,
bias: FloatTensor<Self>,
) -> FloatTensor<Self> {
// Create our zero-sized type that will implement the Backward trait.
#[derive(Debug)]
struct FusedMatmulAddReluBackward;
// Implement the backward trait for the given backend B, the node gradient
// with three other gradients to calculate (lhs, rhs, and bias).
impl<B: Backend> Backward<B, 3> for FusedMatmulAddReluBackward {
// Our state that we must build during the forward pass to compute the backward pass.
//
// Note that we could improve the performance further by only keeping the state of
// tensors that are tracked, improving memory management, but for simplicity, we avoid
// that part.
type State = (NodeId, NodeId, FloatTensor<B>, Shape);
fn backward(
self,
ops: Ops<Self::State, 3>,
grads: &mut Gradients,
checkpointer: &mut Checkpointer,
) {
// Get the nodes of each variable.
let [node_lhs, node_rhs, node_bias] = ops.parents;
// Fetch the gradient for the current node.
let grad = grads.consume::<B>(&ops.node);
// Set our state.
let (lhs_state, rhs_state, output, shape_bias) = ops.state;
let lhs: FloatTensor<B> = checkpointer.retrieve_node_output(lhs_state);
let rhs: FloatTensor<B> = checkpointer.retrieve_node_output(rhs_state);
// Fetch shapes of our tensor to support broadcasting.
let shape_lhs = lhs.shape();
let shape_rhs = rhs.shape();
// Compute the gradient of the output using the already existing `relu_backward`
// function in the basic Burn backend trait.
let grad_output = B::relu_backward(output, grad);
// Compute the lhs gradient, which is the derivative of matmul with support for
// broadcasting.
let grad_lhs = broadcast_shape::<B>(
B::float_matmul(grad_output.clone(), B::float_transpose(rhs)),
&shape_lhs,
);
// Compute the rhs gradient, which is the derivative of matmul with support for
// broadcasting.
let grad_rhs = broadcast_shape::<B>(
B::float_matmul(B::float_transpose(lhs), grad_output.clone()),
&shape_rhs,
);
// The add derivative is only 1, so we just need to support broadcasting to
// compute the bias gradient.
let grad_bias = broadcast_shape::<B>(grad_output, &shape_bias);
// Register the gradient for each variable based on whether they are marked as
// `tracked`.
if let Some(node) = node_bias {
grads.register::<B>(node.id, grad_bias);
}
if let Some(node) = node_lhs {
grads.register::<B>(node.id, grad_lhs);
}
if let Some(node) = node_rhs {
grads.register::<B>(node.id, grad_rhs);
}
}
}
// Prepare a stateful operation with each variable node and corresponding graph.
//
// Each node can be fetched with `ops.parents` in the same order as defined here.
match FusedMatmulAddReluBackward
.prepare::<C>([lhs.node.clone(), rhs.node.clone(), bias.node.clone()])
// Marks the operation as compute bound, meaning it will save its
// state instead of recomputing itself during checkpointing
.compute_bound()
.stateful()
{
OpsKind::Tracked(mut prep) => {
// When at least one node is tracked, we should register our backward step.
// The state consists of what will be needed for this operation's backward pass.
// Since we need the parents' outputs, we must checkpoint their ids to retrieve
// their node output at the beginning of the backward pass. We can also save
// utility data such as the bias shape. If we also need this operation's output,
// we can either save it in the state or recompute it.
// during the backward pass. Here we choose to save it in the state because it's a
// compute bound operation.
let lhs_state = prep.checkpoint(&lhs);
let rhs_state = prep.checkpoint(&rhs);
let bias_shape = bias.primitive.shape();
let output = B::fused_matmul_add_relu(
lhs.primitive.clone(),
rhs.primitive.clone(),
bias.primitive,
);
let state = (lhs_state, rhs_state, output.clone(), bias_shape);
prep.finish(state, output)
}
OpsKind::UnTracked(prep) => {
// When no node is tracked, we can just compute the original operation without
// keeping any state.
let output = B::fused_matmul_add_relu(lhs.primitive, rhs.primitive, bias.primitive);
prep.finish(output)
}
}
}
}
```
The previous code is self-documented to make it clearer, but here is what it does in summary:
We define `fused_matmul_add_relu` within `Autodiff<B>`, allowing any autodiff-decorated backend to
benefit from our implementation. In an autodiff-decorated backend, the forward pass must still be
implemented. This is achieved using a comprehensive match statement block where computation is
delegated to the inner backend, while keeping track of a state. The state comprises any information
relevant to the backward pass, such as input and output tensors, along with the bias shape. When an
operation isn't tracked (meaning there won't be a backward pass for this specific operation in the
graph), storing a state becomes unnecessary, and we simply perform the forward computation.
The backward pass uses the gradient obtained from the preceding node in the computation graph. It
calculates the derivatives for `relu` (`relu_backward`), add (no operation is required here, as the
derivative is one), and `matmul` (another `matmul` with transposed inputs). This results in
gradients for both input tensors and the bias, which are registered for consumption by subsequent
operation nodes.
The only remaining part is to implement our autodiff-decorated backend trait for our JIT Backend.
```rust, ignore
impl<R: CubeRuntime> AutodiffBackend for Autodiff<CubeBackend<R>>
{
}
```
## Conclusion
In this guide, we've implemented a fused kernel using the `cubecl` compiler frontend, enabling
execution on any GPU and any `cubecl` backend. By delving into the inner workings of both the JIT
backend and the autodiff backend, we've gained a deeper understanding of these systems.
While extending a backend may be harder than working with straightforward tensors, the benefits can
be worth it. This approach enables the crafting of custom models with greater control over
execution, which can potentially greatly enhance the performance of your models.
As we conclude this guide, we hope that you have gained insights into Burn's world of backend
extensions, and that it will help you to unleash the full potential of your projects.
burn-book/src/advanced/backend-extension/custom-wgpu-kernel.md
+459
View File
@@ -0,0 +1,459 @@
# Custom WGPU Kernel
In this section, you will learn how to create your own custom operation by writing your own kernel
with the WGPU backend. We will take the example of a common workflow in the deep learning field,
where we create a kernel to fuse multiple operations together. Note that `burn` does this
automatically, but a manual implementation might be more efficient in some cases. We will fuse a
matmul kernel followed by an addition and the ReLU activation function, which is commonly found in
various models. All the code can be found under the
[examples directory](https://github.com/tracel-ai/burn/tree/main/examples/custom-wgpu-kernel).
## Custom Backend Trait
First, we need to determine the type signature of our newly created operation by defining our custom
backend traits. As we will use the associated type `TensorPrimitive` of the `Backend` trait, which
encapsulates the underlying tensor implementation of the backend, we will use a type alias to avoid
the ugly disambiguation with associated types.
```rust, ignore
/// We create our own Backend trait that extends the Burn backend trait.
pub trait Backend: burn::tensor::backend::Backend {
fn fused_matmul_add_relu(
lhs: FloatTensor<Self>,
rhs: FloatTensor<Self>,
bias: FloatTensor<Self>,
) -> FloatTensor<Self>;
}
/// We create our own AutodiffBackend trait that extends the Burn autodiff backend trait.
pub trait AutodiffBackend: Backend + burn::tensor::backend::AutodiffBackend {}
```
In our project, we can use these traits instead of the
`burn::tensor::backend::{Backend, AutodiffBackend}` traits provided by Burn. Burn's user APIs
typically make use of the `Tensor` struct rather than dealing directly with primitive tensor types.
Therefore, we can encapsulate our newly defined backend traits with functions that expose new
operations while maintaining a consistent API.
```rust, ignore
/// We define our custom implementation using the added function on our custom backend.
pub fn matmul_add_relu_custom<B: Backend>(
lhs: Tensor<B, 3>,
rhs: Tensor<B, 3>,
bias: Tensor<B, 3>,
) -> Tensor<B, 3> {
let output = B::fused_matmul_add_relu(
lhs.into_primitive().tensor(),
rhs.into_primitive().tensor(),
bias.into_primitive().tensor(),
);
Tensor::from_primitive(TensorPrimitive::Float(output))
}
/// We define a reference implementation using basic tensor operations.
pub fn matmul_add_relu_reference<B: Backend>(
lhs: Tensor<B, 3>,
rhs: Tensor<B, 3>,
bias: Tensor<B, 3>,
) -> Tensor<B, 3> {
let x = lhs.matmul(rhs) + bias;
activation::relu(x)
}
```
Note that we also provide a reference implementation for testing purposes, which allows us to easily
validate our new implementation. While not mandatory, having a reference implementation can be
valuable, especially in projects where creating a reference implementation solely using basic tensor
operations is feasible.
## Forward Kernel
Now, let's proceed to write the fused kernel using the WGSL shading language. To keep things simple,
we'll create a straightforward matmul kernel without employing any intricate techniques. Although we
won't delve into the details of the WGSL syntax, as it falls beyond the scope of this guide, we
still provide the implementation below for readers who are curious. The actual matmul, add and relu
computations are found at the end, after an extensive overhead whose use is to correctly map each
compute unit to the data it is responsible of, with support for batches.
```wgsl, ignore
@group(0)
@binding(0)
var<storage, read_write> lhs: array<{{ elem }}>;
@group(0)
@binding(1)
var<storage, read_write> rhs: array<{{ elem }}>;
@group(0)
@binding(2)
var<storage, read_write> bias: array<{{ elem }}>;
@group(0)
@binding(3)
var<storage, read_write> output: array<{{ elem }}>;
@group(0)
@binding(4)
var<storage, read_write> info: array<u32>;
const BLOCK_SIZE = {{ workgroup_size_x }}u;
@compute
@workgroup_size({{ workgroup_size_x }}, {{ workgroup_size_y }}, 1)
fn main(
@builtin(global_invocation_id) global_id: vec3<u32>,
@builtin(local_invocation_index) local_idx: u32,
@builtin(workgroup_id) workgroup_id: vec3<u32>,
) {
// Indices
let row = workgroup_id.x * BLOCK_SIZE + (local_idx / BLOCK_SIZE);
let col = workgroup_id.y * BLOCK_SIZE + (local_idx % BLOCK_SIZE);
let batch = global_id.z;
// Basic information
let dim = info[0];
let n_rows = info[6u * dim - 1u];
let n_cols = info[6u * dim];
let K = info[5u * dim - 1u];
// Returns if outside the output dimension
if row >= n_rows || col >= n_cols {
return;
}
// Calculate the corresponding offsets with support for broadcasting.
let offset_output = batch * n_rows * n_cols;
var offset_lhs: u32 = 0u;
var offset_rhs: u32 = 0u;
let batch_dims = dim - 2u;
for (var b: u32 = 1u; b <= batch_dims; b++) {
let stride_lhs = info[b];
let stride_rhs = info[b + dim];
let stride_output = info[b + 2u * dim];
let shape_lhs = info[b + 3u * dim];
let shape_rhs = info[b + 4u * dim];
offset_lhs += offset_output / stride_output % shape_lhs * stride_lhs;
offset_rhs += offset_output / stride_output % shape_rhs * stride_rhs;
}
// Basic matmul implementation
var sum = 0.0;
for (var k: u32 = 0u; k < K; k++) {
let lhs_index = row * K + k;
let rhs_index = k * n_cols + col;
sum += lhs[offset_lhs + lhs_index] * rhs[offset_rhs + rhs_index];
}
let output_index = row * n_cols + col;
let index = offset_output + output_index;
// Add and ReLU
output[index] = max(sum + bias[index], 0.0);
}
```
Now, let's move on to the next step, which involves implementing the remaining code to launch the
kernel. The initial part entails loading the template and populating it with the appropriate
variables. The `register(name, value)` method simply replaces occurrences of `{{ name }}` in the
above WGSL code with some other string before it is compiled. In order to use templating utilities,
you will have to activate the `template` feature of Burn in your `cargo.toml`.
```rust, ignore
// Source the kernel written in WGSL.
kernel_wgsl!(FusedMatmulAddReluRaw, "./kernel.wgsl");
// Define our kernel type with cube information.
#[derive(new, Debug)]
struct FusedMatmulAddRelu<E: FloatElement> {
cube_dim: CubeDim,
_elem: PhantomData<E>,
}
// Implement the dynamic kernel trait for our kernel type.
impl<E: FloatElement> KernelSource for FusedMatmulAddRelu<E> {
fn source(&self) -> SourceTemplate {
// Extend our raw kernel with cube size information using the
// `SourceTemplate` trait.
FusedMatmulAddReluRaw::new()
.source()
.register("workgroup_size_x", self.cube_dim.x.to_string())
.register("workgroup_size_y", self.cube_dim.y.to_string())
.register("elem", E::type_name())
.register("int", "i32")
}
fn id(&self) -> cubecl::KernelId {
cubecl::KernelId::new::<Self>().info(self.cube_dim)
}
}
```
Subsequently, we'll go into implementing our custom backend trait for the WGPU backend. Note that we
won't go into supporting the `fusion` feature flag in this tutorial, so we implement the trait for
the raw `WgpuBackend` type.
```rust, ignore
/// Implement our custom backend trait for the existing backend `WgpuBackend`.
impl Backend for CubeBackend<WgpuRuntime>
{
fn fused_matmul_add_relu(
lhs: FloatTensor<Self>,
rhs: FloatTensor<Self>,
bias: FloatTensor<Self>,
) -> FloatTensor<Self> {
// Define cube dim, hardcoded for simplicity.
let cube_dim = CubeDim { x: 16, y: 16, z: 1 };
lhs.assert_is_on_same_device(&rhs);
lhs.assert_is_on_same_device(&bias);
// For simplicity, make sure each tensor is continuous.
let lhs = into_contiguous(lhs);
let rhs = into_contiguous(rhs);
let bias = into_contiguous(bias);
// Get the matmul relevant shapes.
let ndims = lhs.shape.num_dims();
let num_rows = lhs.shape[ndims - 2];
let num_cols = rhs.shape[ndims - 1];
// Compute shape of output, while tracking number of batches.
let mut num_batches = 1;
let mut shape_out = vec![0; ndims];
for i in shape_out.clone().into_iter().take(ndims - 2) {
shape_out[i] = usize::max(lhs.shape[i], rhs.shape[i]);
num_batches *= shape_out[i];
}
shape_out[ndims - 2] = num_rows;
shape_out[ndims - 1] = num_cols;
let shape_out = Shape::from(shape_out);
// Create a buffer for the output tensor.
let buffer = lhs
.client
.empty(shape_out.num_elements() * core::mem::size_of::<F>());
// Create the output tensor primitive.
let output = CubeTensor::new_contiguous(
lhs.client.clone(),
lhs.device.clone(),
shape_out,
buffer,
F::dtype(),
);
// Create the kernel.
let kernel = FusedMatmulAddRelu::<F>::new(cube_dim);
// Build info buffer with tensor information needed by the kernel, such as shapes and strides.
let info = build_info::<_, F>(&[&lhs, &rhs, &output]);
let info_handle = lhs.client.create(bytemuck::cast_slice(&info));
// Declare the wgsl workgroup with the number of cubes in x, y and z.
let cubes_needed_in_x = f32::ceil(num_rows as f32 / cube_dim.x as f32) as u32;
let cubes_needed_in_y = f32::ceil(num_cols as f32 / cube_dim.y as f32) as u32;
let cube_count =
CubeCount::Static(cubes_needed_in_x, cubes_needed_in_y, num_batches as u32);
// Execute lazily the kernel with the launch information and the given buffers.
lhs.client.execute(
Box::new(SourceKernel::new(kernel, cube_dim)),
cube_count,
Bindings::new().with_buffers(vec![
lhs.handle.binding(),
rhs.handle.binding(),
bias.handle.binding(),
output.handle.clone().binding(),
info_handle.binding(),
]),
);
// Return the output tensor.
output
}
}
```
In the preceding code block, we demonstrated how to launch the kernel that modifies the correct
buffer. It's important to note that Rust's mutability safety doesn't apply here; the context has the
capability to execute any mutable operation on any buffer. While this isn't a problem in the
previous scenario where we only modify the newly created output buffer, it is wise to keep this in
mind.
## Backward
Now that the custom backend trait is implemented for the WGPU backend, you can use it to invoke the
`matmul_add_relu_custom` function. However, calculating gradients is not yet possible at this stage.
If your use case does not extend beyond inference, there is no need to implement any of the
following code.
For the backward pass, we will leverage the backend implementation from `burn-autodiff`, which is
actually generic over the backend. Instead of crafting our own WGSL kernel for the backward pass, we
will use our fused kernel only for the forward pass, and compute the gradient using basic
operations.
```rust, ignore
// Implement our custom backend trait for any backend that also implements our custom backend trait.
//
// Note that we could implement the backend trait only for the Wgpu backend instead of any backend that
// also implements our own API. This would allow us to call any function only implemented for Wgpu
// and potentially call a custom kernel crafted only for this task.
impl<B: Backend, C: CheckpointStrategy> Backend for Autodiff<B, C> {
fn fused_matmul_add_relu(
lhs: FloatTensor<Self>,
rhs: FloatTensor<Self>,
bias: FloatTensor<Self>,
) -> FloatTensor<Self> {
// Create our zero-sized type that will implement the Backward trait.
#[derive(Debug)]
struct FusedMatmulAddReluBackward;
// Implement the backward trait for the given backend B, the node gradient
// with three other gradients to calculate (lhs, rhs, and bias).
impl<B: Backend> Backward<B, 3> for FusedMatmulAddReluBackward {
// Our state that we must build during the forward pass to compute the backward pass.
//
// Note that we could improve the performance further by only keeping the state of
// tensors that are tracked, improving memory management, but for simplicity, we avoid
// that part.
type State = (NodeId, NodeId, FloatTensor<B>, Shape);
fn backward(
self,
ops: Ops<Self::State, 3>,
grads: &mut Gradients,
checkpointer: &mut Checkpointer,
) {
// Get the nodes of each variable.
let [node_lhs, node_rhs, node_bias] = ops.parents;
// Fetch the gradient for the current node.
let grad = grads.consume::<B>(&ops.node);
// Set our state.
let (lhs_state, rhs_state, output, shape_bias) = ops.state;
let lhs: FloatTensor<B> = checkpointer.retrieve_node_output(lhs_state);
let rhs: FloatTensor<B> = checkpointer.retrieve_node_output(rhs_state);
// Fetch shapes of our tensor to support broadcasting.
let shape_lhs = lhs.shape();
let shape_rhs = rhs.shape();
// Compute the gradient of the output using the already existing `relu_backward`
// function in the basic Burn backend trait.
let grad_output = B::relu_backward(output, grad);
// Compute the lhs gradient, which is the derivative of matmul with support for
// broadcasting.
let grad_lhs = broadcast_shape::<B>(
B::float_matmul(grad_output.clone(), B::float_transpose(rhs)),
&shape_lhs,
);
// Compute the rhs gradient, which is the derivative of matmul with support for
// broadcasting.
let grad_rhs = broadcast_shape::<B>(
B::float_matmul(B::float_transpose(lhs), grad_output.clone()),
&shape_rhs,
);
// The add derivative is only 1, so we just need to support broadcasting to
// compute the bias gradient.
let grad_bias = broadcast_shape::<B>(grad_output, &shape_bias);
// Register the gradient for each variable based on whether they are marked as
// `tracked`.
if let Some(node) = node_bias {
grads.register::<B>(node.id, grad_bias);
}
if let Some(node) = node_lhs {
grads.register::<B>(node.id, grad_lhs);
}
if let Some(node) = node_rhs {
grads.register::<B>(node.id, grad_rhs);
}
}
}
// Prepare a stateful operation with each variable node and corresponding graph.
//
// Each node can be fetched with `ops.parents` in the same order as defined here.
match FusedMatmulAddReluBackward
.prepare::<C>([lhs.node.clone(), rhs.node.clone(), bias.node.clone()])
// Marks the operation as compute bound, meaning it will save its
// state instead of recomputing itself during checkpointing
.compute_bound()
.stateful()
{
OpsKind::Tracked(mut prep) => {
// When at least one node is tracked, we should register our backward step.
// The state consists of what will be needed for this operation's backward pass.
// Since we need the parents' outputs, we must checkpoint their ids to retrieve their node
// output at the beginning of the backward. We can also save utility data such as the bias shape
// If we also need this operation's output, we can either save it in the state or recompute it
// during the backward pass. Here we choose to save it in the state because it's a compute bound operation.
let lhs_state = prep.checkpoint(&lhs);
let rhs_state = prep.checkpoint(&rhs);
let bias_shape = bias.primitive.shape();
let output = B::fused_matmul_add_relu(
lhs.primitive.clone(),
rhs.primitive.clone(),
bias.primitive,
);
let state = (lhs_state, rhs_state, output.clone(), bias_shape);
prep.finish(state, output)
}
OpsKind::UnTracked(prep) => {
// When no node is tracked, we can just compute the original operation without
// keeping any state.
let output = B::fused_matmul_add_relu(lhs.primitive, rhs.primitive, bias.primitive);
prep.finish(output)
}
}
}
}
```
The previous code is self-documented to make it clearer, but here is what it does in summary.
We define `fused_matmul_add_relu` within `Autodiff<B>`, allowing any autodiff-decorated backend to
benefit from our implementation. In an autodiff-decorated backend, the forward pass must still be
implemented. This is achieved using a comprehensive match statement block where computation is
delegated to the inner backend, while keeping track of a state. The state comprises any information
relevant to the backward pass, such as input and output tensors, along with the bias shape. When an
operation isn't tracked (meaning there won't be a backward pass for this specific operation in the
graph), storing a state becomes unnecessary, and we simply perform the forward computation.
The backward pass uses the gradient obtained from the preceding node in the computation graph. It
calculates the derivatives for `relu` (`relu_backward`), add (no operation is required here, as the
derivative is one), and `matmul` (another `matmul` with transposed inputs). This results in
gradients for both input tensors and the bias, which are registered for consumption by subsequent
operation nodes.
The only remaining part is to implement our autodiff-decorated backend trait for our WGPU Backend.
```rust, ignore
impl AutodiffBackend for Autodiff<CubeBackend<WgpuRuntime>>
{
}
```
## Conclusion
In this guide, we've implemented a fused kernel using the WGPU backend, enabling execution on any
GPU. By delving into the inner workings of both the WGPU backend and the autodiff backend, we've
gained a deeper understanding of these systems.
While extending a backend may be harder than working with straightforward tensors, the benefits can
be worth it. This approach enables the crafting of custom models with greater control over
execution, which can potentially greatly enhance the performance of your models.
As we conclude this guide, we hope that you have gained insights into Burn's world of backend
extensions, and that it will help you to unleash the full potential of your projects.
burn-book/src/advanced/no-std.md
+100
View File
@@ -0,0 +1,100 @@
# No Standard Library
In this section, you will learn how to run an ONNX inference model on an embedded system, with no
standard library support on a Raspberry Pi Pico 2. This should be universally applicable to other
platforms. All the code can be found in the
[burn-onnx examples](https://github.com/tracel-ai/burn-onnx/tree/main/examples/raspberry-pi-pico).
## Step-by-Step Guide
Let's walk through the process of running an embedded ONNX model:
### Setup
Follow the [embassy guide](https://embassy.dev/book/#_getting_started) for your specific environment. Once setup, you should have something similar to the following.
```
./inference
├── Cargo.lock
├── Cargo.toml
├── build.rs
├── memory.x
└── src
└── main.rs
```
Some other dependencies have to be added
```toml
[dependencies]
embedded-alloc = "0.6.0" # Only if there is no default allocator for your chip
burn = { version = "0.21", default-features = false, features = ["flex"] } # Flex supports no_std
burn-store = { version = "0.21", default-features = false, features = ["burnpack"] }
[build-dependencies]
burn-onnx = { version = "0.21" } # Used to auto generate the rust code to import the model
```
### Import the Model
Follow the directions in [ONNX Import](../onnx-import.md).
Use the following ModelGen config
```rs
ModelGen::new()
.input(my_model)
.out_dir("model/")
.embed_states(true)
.run_from_script();
```
### Global Allocator
First define a global allocator (if you are on a no_std system without alloc).
```rs
use embedded_alloc::LlffHeap as Heap;
#[global_allocator]
static HEAP: Heap = Heap::empty();
#[embassy_executor::main]
async fn main(_spawner: Spawner) {
{
use core::mem::MaybeUninit;
// Watch out for this, if it is too big or small for your model, the
// program may crash. This is in u8 bytes, as such this is a total of 100kb
const HEAP_SIZE: usize = 100 * 1024;
static mut HEAP_MEM: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
unsafe { HEAP.init(&raw mut HEAP_MEM as usize, HEAP_SIZE) } // Initialize the heap
}
}
```
### Define Backend
We are using Flex, so we just need to define the Flex backend as usual
```rs
use burn::{backend::Flex, tensor::{Device, Tensor}};
type Backend = Flex;
type BackendDevice = Device<Backend>;
```
Then inside the `main` function add
```rs
use your_model::Model;
// Get a default device for the backend
let device = BackendDevice::default();
// Create a new model and load the state
let model: Model<Backend> = Model::default();
```
### Running the Model
To run the model, just call it as you would normally
```rs
// Define the tensor
let input = Tensor::<Backend, 2>::from_floats([[input]], &device);
// Run the model on the input
let output = model.forward(input);
```
## Conclusion
Running a model in a no_std environment is pretty much identical to a normal environment. All that is needed is a global allocator.
burn-book/src/advanced/web-assembly.md
+12
View File
@@ -0,0 +1,12 @@
# WebAssembly
Burn supports WebAssembly (WASM) execution using the `Flex` and `WebGpu` backends, allowing
models to run directly in the browser.
Check out the following examples:
- [Image Classification Web](https://github.com/tracel-ai/burn-onnx/tree/main/examples/image-classification-web)
- [MNIST Inference on Web](https://github.com/tracel-ai/burn/tree/main/examples/mnist-inference-web)
When targeting WebAssembly, certain dependencies require additional configuration. In particular,
the `getrandom` crate requires explicit setting when using `WebGpu`.
burn-book/src/basic-workflow/README.md
+33
View File
@@ -0,0 +1,33 @@
# Guide
This guide will walk you through the process of creating a custom model built with Burn. We will
train a simple convolutional neural network model on the MNIST dataset and prepare it for inference.
For clarity, we sometimes omit imports in our code snippets. For more details, please refer to the
corresponding code in the `examples/guide` [directory](https://github.com/tracel-ai/burn/tree/main/examples/guide).
We reproduce this example in a step-by-step fashion, from dataset creation to modeling and training
in the following sections. It is recommended to use the capabilities of your IDE or text editor to
automatically add the missing imports as you add the code snippets to your code.
<div class="warning">
Be sure to checkout the git branch corresponding to the version of Burn you are using to follow
this guide.
The current version of Burn is `0.21` and the corresponding branch to checkout is `main`.
</div>
The code for this demo can be executed from Burn's base directory using the command:
```bash
cargo run --example guide
```
## Key Learnings
- Creating a project
- Creating neural network models
- Importing and preparing datasets
- Training models on data
- Choosing a backend
- Using a model for inference
burn-book/src/basic-workflow/backend.md
+54
View File
@@ -0,0 +1,54 @@
# Backend
We have effectively written most of the necessary code to train our model. However, we have not
explicitly designated the backend to be used at any point. This will be defined in the main
entrypoint of our program, namely the `main` function defined in `src/main.rs`.
```rust , ignore
# #![recursion_limit = "256"]
# mod data;
# mod model;
# mod training;
#
use crate::{model::ModelConfig, training::TrainingConfig};
use burn::{
backend::{Autodiff, Wgpu},
# data::dataset::Dataset,
optim::AdamConfig,
};
fn main() {
type MyBackend = Wgpu<f32, i32>;
type MyAutodiffBackend = Autodiff<MyBackend>;
let device = burn::backend::wgpu::WgpuDevice::default();
let artifact_dir = "/tmp/guide";
crate::training::train::<MyAutodiffBackend>(
artifact_dir,
TrainingConfig::new(ModelConfig::new(10, 512), AdamConfig::new()),
device.clone(),
);
}
```
In this code snippet, we use the `Wgpu` backend which is compatible with any operating system and will
use the GPU. For other options, see the Burn README. This backend type takes the graphics API, the
float type and the int type as generic arguments that will be used during the training. The autodiff
backend is simply the same backend, wrapped within the `Autodiff` struct which imparts differentiability
to any backend.
We call the `train` function defined earlier with a directory for artifacts, the configuration of
the model (the number of digit classes is 10 and the hidden dimension is 512), the optimizer
configuration which in our case will be the default Adam configuration, and the device which can be
obtained from the backend.
You can now train your freshly created model with the command:
```console
cargo run --release
```
When running your project with the command above, you should see the training progression through a
basic CLI dashboard:
<img title="a title" alt="Alt text" src="./training-output.png">
burn-book/src/basic-workflow/data.md
+126
View File
@@ -0,0 +1,126 @@
# Data
Typically, one trains a model on some dataset. Burn provides a library of very useful dataset
sources and transformations, such as Hugging Face dataset utilities that allow to download and store
data into an SQLite database for extremely efficient data streaming and storage. For this guide
though, we will use the MNIST dataset from `burn::data::dataset::vision` which requires no external
dependency.
To iterate over a dataset efficiently, we will define a struct which will implement the `Batcher`
trait. The goal of a batcher is to map individual dataset items into a batched tensor that can be
used as input to our previously defined model.
Let us start by defining our dataset functionalities in a file `src/data.rs`. We shall omit some of
the imports for brevity, but the full code for following this guide can be found at
`examples/guide/` [directory](https://github.com/tracel-ai/burn/tree/main/examples/guide).
```rust , ignore
use burn::{
data::{dataloader::batcher::Batcher, dataset::vision::MnistItem},
prelude::*,
};
#[derive(Clone, Default)]
pub struct MnistBatcher {}
```
This batcher is pretty straightforward, as it only defines a struct that will implement the
`Batcher` trait. The trait is generic over the `Backend` trait, which includes an associated type
for the device, as not all backends expose the same devices. As an example, the Libtorch-based
backend exposes `Cuda(gpu_index)`, `Cpu`, `Vulkan` and `Metal` devices, while the Flex backend
only exposes a single CPU device.
Next, we need to actually implement the batching logic.
```rust , ignore
# use burn::{
# data::{dataloader::batcher::Batcher, dataset::vision::MnistItem},
# prelude::*,
# };
#
# #[derive(Clone, Default)]
# pub struct MnistBatcher {}
#
#[derive(Clone, Debug)]
pub struct MnistBatch<B: Backend> {
pub images: Tensor<B, 3>,
pub targets: Tensor<B, 1, Int>,
}
impl<B: Backend> Batcher<B, MnistItem, MnistBatch<B>> for MnistBatcher {
fn batch(&self, items: Vec<MnistItem>, device: &B::Device) -> MnistBatch<B> {
let images = items
.iter()
.map(|item| TensorData::from(item.image).convert::<B::FloatElem>())
.map(|data| Tensor::<B, 2>::from_data(data, device))
.map(|tensor| tensor.reshape([1, 28, 28]))
// Normalize: scale between [0,1] and make the mean=0 and std=1
// values mean=0.1307,std=0.3081 are from the PyTorch MNIST example
// https://github.com/pytorch/examples/blob/54f4572509891883a947411fd7239237dd2a39c3/mnist/main.py#L122
.map(|tensor| ((tensor / 255) - 0.1307) / 0.3081)
.collect();
let targets = items
.iter()
.map(|item| {
Tensor::<B, 1, Int>::from_data([(item.label as i64).elem::<B::IntElem>()], device)
})
.collect();
let images = Tensor::cat(images, 0);
let targets = Tensor::cat(targets, 0);
MnistBatch { images, targets }
}
}
```
<details>
<summary><strong>🦀 Iterators and Closures</strong></summary>
The iterator pattern allows you to perform some tasks on a sequence of items in turn.
In this example, an iterator is created over the `MnistItem`s in the vector `items` by calling the
`iter` method.
_Iterator adaptors_ are methods defined on the `Iterator` trait that produce different iterators by
changing some aspect of the original iterator. Here, the `map` method is called in a chain to
transform the original data before consuming the final iterator with `collect` to obtain the
`images` and `targets` vectors. Both vectors are then concatenated into a single tensor for the
current batch.
You probably noticed that each call to `map` is different, as it defines a function to execute on
the iterator items at each step. These anonymous functions are called
[_closures_](https://doc.rust-lang.org/book/ch13-01-closures.html) in Rust. They're easy to
recognize due to their syntax which uses vertical bars `||`. The vertical bars capture the input
variables (if applicable) while the rest of the expression defines the function to execute.
If we go back to the example, we can break down and comment the expression used to process the
images.
```rust, ignore
let images = items // take items Vec<MnistItem>
.iter() // create an iterator over it
.map(|item| TensorData::from(item.image).convert::<B::FloatElem>()) // for each item, convert the image to float data struct
.map(|data| Tensor::<B, 2>::from_data(data, device)) // for each data struct, create a tensor on the device
.map(|tensor| tensor.reshape([1, 28, 28])) // for each tensor, reshape to the image dimensions [C, H, W]
.map(|tensor| ((tensor / 255) - 0.1307) / 0.3081) // for each image tensor, apply normalization
.collect(); // consume the resulting iterator & collect the values into a new vector
```
For more information on iterators and closures, be sure to check out the
[corresponding chapter](https://doc.rust-lang.org/book/ch13-00-functional-features.html) in the Rust
Book.
</details><br>
In the previous example, we implement the `Batcher` trait with a list of `MnistItem` as input and a
single `MnistBatch` as output. The batch contains the images in the form of a 3D tensor, along with
a targets tensor that contains the indexes of the correct digit class. The first step is to parse
the image array into a `TensorData` struct. Burn provides the `TensorData` struct to encapsulate
tensor storage information without being specific for a backend. When creating a tensor from data,
we often need to convert the data precision to the current backend in use. This can be done with the
`.convert()` method (in this example, the data is converted backend's float element type
`B::FloatElem`). While importing the `burn::tensor::ElementConversion` trait, you can call `.elem()`
on a specific number to convert it to the current backend element type in use.
burn-book/src/basic-workflow/inference.md
+89
View File
@@ -0,0 +1,89 @@
# Inference
Now that we have trained our model, the next natural step is to use it for inference.
You need two things in order to load weights for a model: the model's record and the model's config.
Since parameters in Burn are lazy initialized, no allocation and GPU/CPU kernels are executed by the
`ModelConfig::init` function. The weights are initialized when used for the first time, therefore
you can safely use `config.init(device).load_record(record)` without any meaningful performance
cost. Let's create a simple `infer` method in a new file `src/inference.rs` which we will use to
load our trained model.
```rust , ignore
# use crate::{data::MnistBatcher, training::TrainingConfig};
# use burn::{
# data::{dataloader::batcher::Batcher, dataset::vision::MnistItem},
# prelude::*,
# store::ModuleRecord,
# };
#
pub fn infer<B: Backend>(artifact_dir: &str, device: B::Device, item: MnistItem) {
let config = TrainingConfig::load(format!("{artifact_dir}/config.json"))
.expect("Config should exist for the model; run train first");
let record = ModuleRecord::load(format!("{artifact_dir}/model"))
.expect("Trained model should exist; run train first");
let model = config.model.init::<B>(&device).load_record(record);
let label = item.label;
let batcher = MnistBatcher::default();
let batch = batcher.batch(vec![item], &device);
let output = model.forward(batch.images);
let predicted = output.argmax(1).flatten::<1>(0, 1).into_scalar();
println!("Predicted {predicted} Expected {label}");
}
```
The first step is to load the configuration of the training to fetch the correct model
configuration. Then we can load the saved record from its burnpack file. Finally we can init the
model with the configuration and apply the record. For simplicity we can use the
same batcher used during the training to pass from a MnistItem to a tensor.
By running the infer function, you should see the predictions of your model!
Add the call to `infer` to the `main.rs` file after the `train` function call:
```rust , ignore
# #![recursion_limit = "256"]
# mod data;
# mod inference;
# mod model;
# mod training;
#
# use crate::{model::ModelConfig, training::TrainingConfig};
# use burn::{
# backend::{Autodiff, Wgpu},
# data::dataset::Dataset,
# optim::AdamConfig,
# };
#
# fn main() {
# type MyBackend = Wgpu<f32, i32>;
# type MyAutodiffBackend = Autodiff<MyBackend>;
#
# let device = burn::backend::wgpu::WgpuDevice::default();
# let artifact_dir = "/tmp/guide";
# crate::training::train::<MyAutodiffBackend>(
# artifact_dir,
# TrainingConfig::new(ModelConfig::new(10, 512), AdamConfig::new()),
# device.clone(),
# );
crate::inference::infer::<MyBackend>(
artifact_dir,
device,
burn::data::dataset::vision::MnistDataset::test()
.get(42)
.unwrap(),
);
# }
```
The number `42` is the index of the image in the MNIST dataset. You can explore and verify them
using this [MNIST viewer](https://observablehq.com/@davidalber/mnist-viewer).
---
In this short guide, we've introduced you to the fundamental building blocks for getting started
with Burn. While there's still plenty to explore, our goal has been to provide you with the
essential knowledge to kickstart your productivity within the framework.
burn-book/src/basic-workflow/model.md
+406
View File
@@ -0,0 +1,406 @@
# Model
The first step is to create a project and add the different Burn dependencies. Start by creating a
new project with Cargo:
```console
cargo new guide
```
As [mentioned previously](../getting-started.md#creating-a-burn-application), this will initialize
your `guide` project directory with a `Cargo.toml` and a `src/main.rs` file.
In the `Cargo.toml` file, add the `burn` dependency with `train`, `vision` and `wgpu` features.
Since we disable the default features, we also want to enable `std`, `tui` (for the dashboard) and
`fusion` for wgpu. Then run `cargo build` to build the project and import all the dependencies.
```toml
[package]
name = "guide"
version = "0.1.0"
edition = "2024"
[dependencies]
# Disable autotune default for convolutions
burn = { version = "~0.21", features = ["std", "tui", "train", "vision", "wgpu", "fusion"], default-features = false }
# burn = { version = "~0.21", features = ["train", "vision", "wgpu"] }
```
Our goal will be to create a basic convolutional neural network used for image classification. We
will keep the model simple by using two convolution layers followed by two linear layers, some
pooling and ReLU activations. We will also use dropout to improve training performance.
Let us start by defining our model struct in a new file `src/model.rs`.
```rust , ignore
use burn::{
nn::{
conv::{Conv2d, Conv2dConfig},
pool::{AdaptiveAvgPool2d, AdaptiveAvgPool2dConfig},
Dropout, DropoutConfig, Linear, LinearConfig, Relu,
},
prelude::*,
};
#[derive(Module, Debug)]
pub struct Model<B: Backend> {
conv1: Conv2d<B>,
conv2: Conv2d<B>,
pool: AdaptiveAvgPool2d,
dropout: Dropout,
linear1: Linear<B>,
linear2: Linear<B>,
activation: Relu,
}
```
There are two major things going on in this code sample.
1. You can create a deep learning module with the `#[derive(Module)]` attribute on top of a struct.
This will generate the necessary code so that the struct implements the `Module` trait. This
trait will make your module both trainable and (de)serializable while adding related
functionalities. Like other attributes often used in Rust, such as `Clone`, `PartialEq` or
`Debug`, each field within the struct must also implement the `Module` trait.
<details>
<summary><strong>🦀 Trait</strong></summary>
Traits are a powerful and flexible Rust language feature. They provide a way to define shared
behavior for a particular type, which can be shared with other types.
A type's behavior consists of the methods called on that type. Since all `Module`s should
implement the same functionality, it is defined as a trait. Implementing a trait on a particular
type usually requires the user to implement the defined behaviors of the trait for their types,
though that is not the case here as explained above with the `derive` attribute. Check out the
[explainer below](#derive-attribute) to learn why.
For more details on traits, take a look at the
[associated chapter](https://doc.rust-lang.org/book/ch10-02-traits.html) in the Rust Book.
</details><br>
<details id="derive-attribute">
<summary><strong>🦀 Derive Macro</strong></summary>
The `derive` attribute allows traits to be implemented easily by generating code that will
implement a trait with its own default implementation on the type that was annotated with the
`derive` syntax.
This is accomplished through a feature of Rust called
[procedural macros](https://doc.rust-lang.org/reference/procedural-macros.html), which allow us
to run code at compile time that operates over Rust syntax, both consuming and producing Rust
syntax. Using the attribute `#[my_macro]`, you can effectively extend the provided code. You will
see that the derive macro is very frequently employed to recursively implement traits, where the
implementation consists of the composition of all fields.
In this example, we want to derive the [`Module`](../building-blocks/module.md) and `Debug`
traits.
```rust, ignore
#[derive(Module, Debug)]
pub struct MyCustomModule<B: Backend> {
linear1: Linear<B>,
linear2: Linear<B>,
activation: Relu,
}
```
The basic `Debug` implementation is provided by the compiler to format a value using the `{:?}`
formatter. For ease of use, the `Module` trait implementation is automatically handled by Burn so
you don't have to do anything special. It essentially acts as parameter container.
For more details on derivable traits, take a look at the Rust
[appendix](https://doc.rust-lang.org/book/appendix-03-derivable-traits.html),
[reference](https://doc.rust-lang.org/reference/attributes/derive.html) or
[example](https://doc.rust-lang.org/rust-by-example/trait/derive.html).
</details><br>
2. Note that the struct is generic over the [`Backend`](../building-blocks/backend.md) trait. The
backend trait abstracts the underlying low level implementations of tensor operations, allowing
your new model to run on any backend. Contrary to other frameworks, the backend abstraction isn't
determined by a compilation flag or a device type. This is important because you can extend the
functionalities of a specific backend (see
[backend extension section](../advanced/backend-extension)), and it allows for an innovative
[autodiff system](../building-blocks/autodiff.md). You can also change backend during runtime,
for instance to compute training metrics on a cpu backend while using a gpu one only to train the
model. In our example, the backend in use will be determined later on.
<details>
<summary><strong>🦀 Trait Bounds</strong></summary>
Trait bounds provide a way for generic items to restrict which types are used as their
parameters. The trait bounds stipulate what functionality a type implements. Therefore, bounding
restricts the generic to types that conform to the bounds. It also allows generic instances to
access the methods of traits specified in the bounds.
For a simple but concrete example, check out the
[Rust By Example on bounds](https://doc.rust-lang.org/rust-by-example/generics/bounds.html).
In Burn, the `Backend` trait enables you to run tensor operations using different implementations
as it abstracts tensor, device and element types. The
[getting started example](../getting-started.md#writing-a-code-snippet) illustrates the advantage
of having a simple API that works for different backend implementations. While it used the WGPU
backend, you could easily swap it with any other supported backend.
```rust, ignore
// Choose from any of the supported backends.
// type Backend = Cuda;
// type Backend = Flex;
type Backend = Wgpu;
// Creation of two tensors.
let tensor_1 = Tensor::<Backend, 2>::from_data([[2., 3.], [4., 5.]], &device);
let tensor_2 = Tensor::<Backend, 2>::ones_like(&tensor_1);
// Print the element-wise addition (done with the selected backend) of the two tensors.
println!("{}", tensor_1 + tensor_2);
```
For more details on trait bounds, check out the Rust
[trait bound section](https://doc.rust-lang.org/book/ch10-02-traits.html#trait-bound-syntax) or
[reference](https://doc.rust-lang.org/reference/items/traits.html#trait-bounds).
</details><br>
Note that each time you create a new file in the `src` directory you also need to explicitly add
this module to the `main.rs` file. For instance after creating the `model.rs`, you need to add the
following at the top of the main file:
```rust , ignore
mod model;
#
# fn main() {
# }
```
Next, we need to instantiate the model for training.
```rust , ignore
# use burn::{
# nn::{
# conv::{Conv2d, Conv2dConfig},
# pool::{AdaptiveAvgPool2d, AdaptiveAvgPool2dConfig},
# Dropout, DropoutConfig, Linear, LinearConfig, Relu,
# },
# prelude::*,
# };
#
# #[derive(Module, Debug)]
# pub struct Model<B: Backend> {
# conv1: Conv2d<B>,
# conv2: Conv2d<B>,
# pool: AdaptiveAvgPool2d,
# dropout: Dropout,
# linear1: Linear<B>,
# linear2: Linear<B>,
# activation: Relu,
# }
#
#[derive(Config, Debug)]
pub struct ModelConfig {
num_classes: usize,
hidden_size: usize,
#[config(default = "0.5")]
dropout: f64,
}
impl ModelConfig {
/// Returns the initialized model.
pub fn init<B: Backend>(&self, device: &B::Device) -> Model<B> {
Model {
conv1: Conv2dConfig::new([1, 8], [3, 3]).init(device),
conv2: Conv2dConfig::new([8, 16], [3, 3]).init(device),
pool: AdaptiveAvgPool2dConfig::new([8, 8]).init(),
activation: Relu::new(),
linear1: LinearConfig::new(16 * 8 * 8, self.hidden_size).init(device),
linear2: LinearConfig::new(self.hidden_size, self.num_classes).init(device),
dropout: DropoutConfig::new(self.dropout).init(),
}
}
}
```
At a glance, you can view the model configuration by printing the model instance:
```rust , ignore
#![recursion_limit = "256"]
mod model;
use crate::model::ModelConfig;
use burn::backend::Wgpu;
fn main() {
type MyBackend = Wgpu<f32, i32>;
let device = Default::default();
let model = ModelConfig::new(10, 512).init::<MyBackend>(&device);
println!("{model}");
}
```
Output:
```rust , ignore
Model {
conv1: Conv2d {ch_in: 1, ch_out: 8, stride: [1, 1], kernel_size: [3, 3], dilation: [1, 1], groups: 1, padding: Valid, params: 80}
conv2: Conv2d {ch_in: 8, ch_out: 16, stride: [1, 1], kernel_size: [3, 3], dilation: [1, 1], groups: 1, padding: Valid, params: 1168}
pool: AdaptiveAvgPool2d {output_size: [8, 8]}
dropout: Dropout {prob: 0.5}
linear1: Linear {d_input: 1024, d_output: 512, bias: true, params: 524800}
linear2: Linear {d_input: 512, d_output: 10, bias: true, params: 5130}
activation: Relu
params: 531178
}
```
<details>
<summary><strong>🦀 References</strong></summary>
In the previous example, the `init()` method signature uses `&` to indicate that the parameter types
are references: `&self`, a reference to the current receiver (`ModelConfig`), and
`device: &B::Device`, a reference to the backend device.
```rust, ignore
pub fn init<B: Backend>(&self, device: &B::Device) -> Model<B> {
Model {
// ...
}
}
```
References in Rust allow us to point to a resource to access its data without owning it. The idea of
ownership is quite core to Rust and is worth
[reading up on](https://doc.rust-lang.org/book/ch04-00-understanding-ownership.html).
In a language like C, memory management is explicit and up to the programmer, which means it is easy
to make mistakes. In a language like Java or Python, memory management is automatic with the help of
a garbage collector. This is very safe and straightforward, but also incurs a runtime cost.
In Rust, memory management is rather unique. Aside from simple types that implement
[`Copy`](https://doc.rust-lang.org/std/marker/trait.Copy.html) (e.g.,
[primitives](https://doc.rust-lang.org/rust-by-example/primitives.html) like integers, floats,
booleans and `char`), every value is _owned_ by some variable called the _owner_. Ownership can be
transferred from one variable to another and sometimes a value can be _borrowed_. Once the _owner_
variable goes out of scope, the value is _dropped_, which means that any memory it allocated can be
freed safely.
Because the method does not own the `self` and `device` variables, the values the references point
to will not be dropped when the reference stops being used (i.e., the scope of the method).
For more information on references and borrowing, be sure to read the
[corresponding chapter](https://doc.rust-lang.org/book/ch04-02-references-and-borrowing.html) in the
Rust Book.
</details><br>
When creating a custom neural network module, it is often a good idea to create a config alongside
the model struct. This allows you to define default values for your network, thanks to the `Config`
attribute. The benefit of this attribute is that it makes the configuration serializable, enabling
you to painlessly save your model hyperparameters, enhancing your experimentation process. Note that
a constructor will automatically be generated for your configuration, which will take in as input
values the parameters which do not have default values:
`let config = ModelConfig::new(num_classes, hidden_size);`. The default values can be overridden
easily with builder-like methods: (e.g `config.with_dropout(0.2);`)
The first implementation block is related to the initialization method. As we can see, all fields
are set using the configuration of the corresponding neural network's underlying module. In this
specific case, we have chosen to expand the tensor channels from 1 to 8 with the first layer, then
from 8 to 16 with the second layer, using a kernel size of 3 on all dimensions. We also use the
adaptive average pooling module to reduce the dimensionality of the images to an 8 by 8 matrix,
which we will flatten in the forward pass to have a 1024 (16 * 8 * 8) resulting tensor.
Now let's see how the forward pass is defined.
```rust , ignore
# use burn::{
# nn::{
# conv::{Conv2d, Conv2dConfig},
# pool::{AdaptiveAvgPool2d, AdaptiveAvgPool2dConfig},
# Dropout, DropoutConfig, Linear, LinearConfig, Relu,
# },
# prelude::*,
# };
#
# #[derive(Module, Debug)]
# pub struct Model<B: Backend> {
# conv1: Conv2d<B>,
# conv2: Conv2d<B>,
# pool: AdaptiveAvgPool2d,
# dropout: Dropout,
# linear1: Linear<B>,
# linear2: Linear<B>,
# activation: Relu,
# }
#
# #[derive(Config, Debug)]
# pub struct ModelConfig {
# num_classes: usize,
# hidden_size: usize,
# #[config(default = "0.5")]
# dropout: f64,
# }
#
# impl ModelConfig {
# /// Returns the initialized model.
# pub fn init<B: Backend>(&self, device: &B::Device) -> Model<B> {
# Model {
# conv1: Conv2dConfig::new([1, 8], [3, 3]).init(device),
# conv2: Conv2dConfig::new([8, 16], [3, 3]).init(device),
# pool: AdaptiveAvgPool2dConfig::new([8, 8]).init(),
# activation: Relu::new(),
# linear1: LinearConfig::new(16 * 8 * 8, self.hidden_size).init(device),
# linear2: LinearConfig::new(self.hidden_size, self.num_classes).init(device),
# dropout: DropoutConfig::new(self.dropout).init(),
# }
# }
# }
#
impl<B: Backend> Model<B> {
/// # Shapes
/// - Images [batch_size, height, width]
/// - Output [batch_size, num_classes]
pub fn forward(&self, images: Tensor<B, 3>) -> Tensor<B, 2> {
let [batch_size, height, width] = images.dims();
// Create a channel at the second dimension.
let x = images.reshape([batch_size, 1, height, width]);
let x = self.conv1.forward(x); // [batch_size, 8, _, _]
let x = self.dropout.forward(x);
let x = self.conv2.forward(x); // [batch_size, 16, _, _]
let x = self.dropout.forward(x);
let x = self.activation.forward(x);
let x = self.pool.forward(x); // [batch_size, 16, 8, 8]
let x = x.reshape([batch_size, 16 * 8 * 8]);
let x = self.linear1.forward(x);
let x = self.dropout.forward(x);
let x = self.activation.forward(x);
self.linear2.forward(x) // [batch_size, num_classes]
}
}
```
For former PyTorch users, this might feel very intuitive, as each module is directly incorporated
into the code using an eager API. Note that no abstraction is imposed for the forward method. You
are free to define multiple forward functions with the names of your liking. Most of the neural
network modules already built with Burn use the `forward` nomenclature, simply because it is
standard in the field.
Similar to neural network modules, the [`Tensor`](../building-blocks/tensor.md) struct given as a
parameter also takes the Backend trait as a generic argument, alongside its dimensionality. Even if
it is not used in this specific example, it is possible to add the kind of the tensor as a third
generic argument. For example, a 3-dimensional Tensor of different data types(float, int, bool)
would be defined as following:
```rust , ignore
Tensor<B, 3> // Float tensor (default)
Tensor<B, 3, Float> // Float tensor (explicit)
Tensor<B, 3, Int> // Int tensor
Tensor<B, 3, Bool> // Bool tensor
```
Note that the specific element type, such as `f16`, `f32` and the likes, will be defined later with
the backend.
burn-book/src/basic-workflow/training-output.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 209 KiB

burn-book/src/basic-workflow/training.md
+297
View File
@@ -0,0 +1,297 @@
# Training
We are now ready to write the necessary code to train our model on the MNIST dataset. We shall
define the code for this training section in the file: `src/training.rs`.
Instead of a simple tensor, the model should output an item that can be understood by the learner, a
struct whose responsibility is to apply an optimizer to the model. The output struct is used for all
metrics calculated during the training. Therefore it should include all the necessary information to
calculate any metric that you want for a task.
Burn provides two basic output types: `ClassificationOutput` and `RegressionOutput`. They implement
the necessary trait to be used with metrics. It is possible to create your own item, but it is
beyond the scope of this guide.
Since the MNIST task is a classification problem, we will use the `ClassificationOutput` type.
```rust , ignore
# use crate::{
# data::{MnistBatch, MnistBatcher},
# model::{Model, ModelConfig},
# };
# use burn::{
# data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset},
# nn::loss::CrossEntropyLossConfig,
# optim::AdamConfig,
# prelude::*,
# tensor::backend::AutodiffBackend,
# train::{
# ClassificationOutput, Learner, SupervisedTraining, TrainOutput, TrainStep, InferenceStep,
# metric::{AccuracyMetric, LossMetric},
# },
# };
#
impl<B: Backend> Model<B> {
pub fn forward_classification(
&self,
images: Tensor<B, 3>,
targets: Tensor<B, 1, Int>,
) -> ClassificationOutput<B> {
let output = self.forward(images);
let loss = CrossEntropyLossConfig::new()
.init(&output.device())
.forward(output.clone(), targets.clone());
ClassificationOutput::new(loss, output, targets)
}
}
```
As evident from the preceding code block, we employ the cross-entropy loss module for loss
calculation, without the inclusion of any padding token. We then return the classification output
containing the loss, the output tensor with all logits and the targets.
Please take note that tensor operations receive owned tensors as input. For reusing a tensor
multiple times, you need to use the `clone()` function. There's no need to worry; this process won't
involve actual copying of the tensor data. Instead, it will simply indicate that the tensor is
employed in multiple instances, implying that certain operations won't be performed in place. In
summary, our API has been designed with owned tensors to optimize performance.
Moving forward, we will proceed with the implementation of both the training and validation steps
for our model.
```rust , ignore
# use crate::{
# data::{MnistBatch, MnistBatcher},
# model::{Model, ModelConfig},
# };
# use burn::{
# data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset},
# nn::loss::CrossEntropyLossConfig,
# optim::AdamConfig,
# prelude::*,
# tensor::backend::AutodiffBackend,
# train::{
# ClassificationOutput, InferenceStep, Learner, SupervisedTraining, TrainOutput, TrainStep,
# metric::{AccuracyMetric, LossMetric},
# },
# };
#
# impl<B: Backend> Model<B> {
# pub fn forward_classification(
# &self,
# images: Tensor<B, 3>,
# targets: Tensor<B, 1, Int>,
# ) -> ClassificationOutput<B> {
# let output = self.forward(images);
# let loss = CrossEntropyLossConfig::new()
# .init(&output.device())
# .forward(output.clone(), targets.clone());
#
# ClassificationOutput::new(loss, output, targets)
# }
# }
impl<B: AutodiffBackend> TrainStep for Model<B> {
type Input = MnistBatch<B>;
type Output = ClassificationOutput<B>;
fn step(&self, batch: MnistBatch<B>) -> TrainOutput<ClassificationOutput<B>> {
let item = self.forward_classification(batch.images, batch.targets);
TrainOutput::new(self, item.loss.backward(), item)
}
}
impl<B: Backend> InferenceStep for Model<B> {
type Input = MnistBatch<B>;
type Output = ClassificationOutput<B>;
fn step(&self, batch: MnistBatch<B>) -> ClassificationOutput<B> {
self.forward_classification(batch.images, batch.targets)
}
}
```
Here we define the input and output types as generic arguments in the `TrainStep` and `InferenceStep`.
We will call them `MnistBatch` and `ClassificationOutput`. In the training step, the computation of
gradients is straightforward, necessitating a simple invocation of `backward()` on the loss. Note
that contrary to PyTorch, gradients are not stored alongside each tensor parameter, but are rather
returned by the backward pass, as such: `let gradients = loss.backward();`. The gradient of a
parameter can be obtained with the grad function: `let grad = tensor.grad(&gradients);`. Although it
is not necessary when using the learner struct and the optimizers, it can prove to be quite useful
when debugging or writing custom training loops. One of the differences between the training and the
validation steps is that the former requires the backend to implement `AutodiffBackend` and not just
`Backend`. Otherwise, the `backward` function is not available, as the backend does not support
autodiff. We will see later how to create a backend with autodiff support.
<details>
<summary><strong>🦀 Generic Type Constraints in Method Definitions</strong></summary>
Although generic data types, trait and trait bounds were already introduced in previous sections of
this guide, the previous code snippet might be a lot to take in at first.
In the example above, we implement the `TrainStep` and `InferenceStep` trait for our `Model` struct,
which is generic over the `Backend` trait as has been covered before. These traits are provided by
`burn::train` and define a common `step` method that should be implemented for all structs. Since
the trait is generic over the input and output types, the trait implementation must specify the
concrete types used. This is where the additional type constraints appear
`<MnistBatch<B>, ClassificationOutput<B>>`. As we saw previously, the concrete input type for the
batch is `MnistBatch`, and the output of the forward pass is `ClassificationOutput`. The `step`
method signature matches the concrete input and output types.
For more details specific to constraints on generic types when defining methods, take a look at
[this section](https://doc.rust-lang.org/book/ch10-01-syntax.html#in-method-definitions) of the Rust
Book.
</details><br>
Let us move on to establishing the practical training configuration.
```rust , ignore
# use std::path::PathBuf;
#
# use crate::{
# data::{MnistBatch, MnistBatcher},
# model::{Model, ModelConfig},
# };
# use burn::{
# data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset},
# nn::loss::CrossEntropyLossConfig,
# optim::AdamConfig,
# prelude::*,
# tensor::backend::AutodiffBackend,
# train::{
# ClassificationOutput, InferenceStep, Learner, SupervisedTraining, TrainOutput, TrainStep,
# metric::{AccuracyMetric, LossMetric},
# },
# };
#
# impl<B: Backend> Model<B> {
# pub fn forward_classification(
# &self,
# images: Tensor<B, 3>,
# targets: Tensor<B, 1, Int>,
# ) -> ClassificationOutput<B> {
# let output = self.forward(images);
# let loss = CrossEntropyLossConfig::new()
# .init(&output.device())
# .forward(output.clone(), targets.clone());
#
# ClassificationOutput::new(loss, output, targets)
# }
# }
# impl<B: AutodiffBackend> TrainStep for Model<B> {
# type Input = MnistBatch<B>;
# type Output = ClassificationOutput<B>;
#
# fn step(&self, batch: MnistBatch<B>) -> TrainOutput<ClassificationOutput<B>> {
# let item = self.forward_classification(batch.images, batch.targets);
#
# TrainOutput::new(self, item.loss.backward(), item)
# }
# }
#
# impl<B: Backend> InferenceStep for Model<B> {
# type Input = MnistBatch<B>;
# type Output = ClassificationOutput<B>;
#
# fn step(&self, batch: MnistBatch<B>) -> ClassificationOutput<B> {
# self.forward_classification(batch.images, batch.targets)
# }
# }
#
#[derive(Config, Debug)]
pub struct TrainingConfig {
pub model: ModelConfig,
pub optimizer: AdamConfig,
#[config(default = 10)]
pub num_epochs: usize,
#[config(default = 64)]
pub batch_size: usize,
#[config(default = 4)]
pub num_workers: usize,
#[config(default = 42)]
pub seed: u64,
#[config(default = 1.0e-4)]
pub learning_rate: f64,
}
fn create_artifact_dir(artifact_dir: &str) {
std::fs::remove_file(PathBuf::from(artifact_dir).join("experiment.log")).ok();
std::fs::create_dir_all(artifact_dir).ok();
}
pub fn train<B: AutodiffBackend>(artifact_dir: &str, config: TrainingConfig, device: B::Device) {
create_artifact_dir(artifact_dir);
config
.save(format!("{artifact_dir}/config.json"))
.expect("Config should be saved successfully");
B::seed(&device, config.seed);
let batcher = MnistBatcher::default();
let dataloader_train = DataLoaderBuilder::new(batcher.clone())
.batch_size(config.batch_size)
.shuffle(config.seed)
.num_workers(config.num_workers)
.build(MnistDataset::train());
let dataloader_test = DataLoaderBuilder::new(batcher)
.batch_size(config.batch_size)
.shuffle(config.seed)
.num_workers(config.num_workers)
.build(MnistDataset::test());
let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test)
.metrics((AccuracyMetric::new(), LossMetric::new()))
.with_checkpointer()
.num_epochs(config.num_epochs)
.summary();
let model = config.model.init::<B>(&device);
let result = training.launch(Learner::new(
model,
config.optimizer.init(),
config.learning_rate,
));
result
.model
.into_record()
.save(format!("{artifact_dir}/model"))
.expect("Trained model should be saved successfully");
}
```
It is a good practice to use the `Config` derive to create the experiment configuration. In the
`train` function, the first thing we are doing is making sure the `artifact_dir` exists, using the
standard rust library for file manipulation. All checkpoints, logging and metrics will be stored
under this directory. We initialize the dataloaders using the previously created batcher. Since no
automatic differentiation is needed during the validation phase, the `training.launch(...)` method
defines the necessary backend bounds on the data loader for `B::InnerBackend` (see
[Backend](./backend.md)). The autodiff capabilities are available through a type system, making it
nearly impossible to forget to deactivate gradient calculation.
Next, we create a supervised training runner with the dataloaders for training and validation and
we register the accuracy and loss metric on both training and validation steps. We also enable
checkpointing with `with_checkpointer()`, which periodically saves the model, optimizer, and learning
rate scheduler state to burnpack files under the experiment directory so training can be resumed.
For the sake of simplicity in this example, we employ the test set as the validation
set; however, we do not recommend this practice for actual usage.
We create the learner containing the model, the optimizer and the learning rate. Notably, the third
argument of the learner's `new` function should actually be a learning rate _scheduler_. When provided with a
float as in our example, it is automatically transformed into a _constant_ learning rate scheduler.
The learning rate is not part of the optimizer config as it is often done in other frameworks, but
rather passed as a parameter when executing the optimizer step. This avoids having to mutate the
state of the optimizer and is therefore more functional. It makes no difference when using the
learner struct, but it will be an essential nuance to grasp if you implement your own training loop.
Once the learner and supervised training instance are created, we can call `training.launch` and provide the learner.
Finally, the trained model is returned by the `launch` method. The trained weights are then saved by
taking a record with `into_record()` and calling `save`, which writes a burnpack (`.bpk`) file. A
record holds plain tensor data, so any backend, regardless of precision, can load recorded weights of
any kind.
burn-book/src/building-blocks/README.md
+7
View File
@@ -0,0 +1,7 @@
# Building Blocks
In this section, we'll guide you through the core elements that make up Burn. We'll walk you through
the key components that serve as the building blocks of the framework and your future projects.
As you explore Burn, you might notice that we occasionally draw comparisons to PyTorch. We believe
it can provide a smoother learning curve and help you grasp the nuances more effectively.
burn-book/src/building-blocks/autodiff.md
+90
View File
@@ -0,0 +1,90 @@
# Autodiff
Burn's tensor also supports auto-differentiation, which is an essential part of any deep learning
framework. We introduced the `Backend` trait in the [previous section](./backend.md), but Burn also
has another trait for autodiff: `AutodiffBackend`.
However, not all tensors support auto-differentiation; you need a backend that implements both the
`Backend` and `AutodiffBackend` traits. Fortunately, you can add auto-differentiation capabilities to any
backend using a backend decorator: `type MyAutodiffBackend = Autodiff<MyBackend>`. This
decorator implements both the `AutodiffBackend` and `Backend` traits by maintaining a dynamic
computational graph and utilizing the inner backend to execute tensor operations.
The `AutodiffBackend` trait adds new operations on float tensors that can't be called otherwise. It also
provides a new associated type, `B::Gradients`, where each calculated gradient resides.
```rust, ignore
fn calculate_gradients<B: AutodiffBackend>(tensor: Tensor<B, 2>) -> B::Gradients {
let mut gradients = tensor.clone().backward();
let tensor_grad = tensor.grad(&gradients); // get
let tensor_grad = tensor.grad_remove(&mut gradients); // pop
gradients
}
```
Note that some functions will always be available even if the backend doesn't implement the
`AutodiffBackend` trait. In such cases, those functions will do nothing.
| Burn API | PyTorch Equivalent |
| --------------------------------------- | ----------------------------- |
| `tensor.detach()` | `tensor.detach()` |
| `tensor.require_grad()` | `tensor.requires_grad()` |
| `tensor.is_require_grad()` | `tensor.requires_grad` |
| `tensor.set_require_grad(require_grad)` | `tensor.requires_grad(False)` |
However, you're unlikely to make any mistakes since you can't call `backward` on a tensor that is on
a backend that doesn't implement `AutodiffBackend`. Additionally, you can't retrieve the gradient of a
tensor without an autodiff backend.
## Difference with PyTorch
The way Burn handles gradients is different from PyTorch. First, when calling `backward`, each
parameter doesn't have its `grad` field updated. Instead, the backward pass returns all the
calculated gradients in a container. This approach offers numerous benefits, such as the ability to
easily send gradients to other threads.
You can also retrieve the gradient for a specific parameter using the `grad` method on a tensor.
Since this method takes the gradients as input, it's hard to forget to call `backward` beforehand.
Note that sometimes, using `grad_remove` can improve performance by allowing inplace operations.
In PyTorch, when you don't need gradients for inference or validation, you typically need to scope
your code using a block.
```python
# Inference mode
torch.inference():
# your code
...
# Or no grad
torch.no_grad():
# your code
...
```
With Burn, you don't need to wrap the backend with the `Autodiff` for inference, and you
can call `inner()` to obtain the inner tensor, which is useful for validation.
```rust, ignore
/// Use `B: AutodiffBackend`
fn example_validation<B: AutodiffBackend>(tensor: Tensor<B, 2>) {
let inner_tensor: Tensor<B::InnerBackend, 2> = tensor.inner();
let _ = inner_tensor + 5;
}
/// Use `B: Backend`
fn example_inference<B: Backend>(tensor: Tensor<B, 2>) {
let _ = tensor + 5;
...
}
```
**Gradients with Optimizers**
We've seen how gradients can be used with tensors, but the process is a bit different when working
with optimizers from `burn-core`. To work with the `Module` trait, a translation step is required to
link tensor parameters with their gradients. This step is necessary to easily support gradient
accumulation and training on multiple devices, where each module can be forked and run on different
devices in parallel. We'll explore deeper into this topic in the [Module](./module.md) section.
burn-book/src/building-blocks/backend.md
+14
View File
@@ -0,0 +1,14 @@
# Backend
Nearly everything in Burn is based on the `Backend` trait, which enables you to run tensor
operations using different implementations without having to modify your code. While a backend may
not necessarily have autodiff capabilities, the `AutodiffBackend` trait specifies when autodiff is
needed. This trait not only abstracts operations but also tensor, device, and element types,
providing each backend the flexibility they need. It's worth noting that the trait assumes eager
mode since burn fully supports dynamic graphs. However, we may create another API to assist with
integrating graph-based backends, without requiring any changes to the user's code.
Users are not expected to directly use the backend trait methods, as it is primarily designed with
backend developers in mind rather than Burn users. Therefore, most Burn userland APIs are generic
across backends. This approach helps users discover the API more organically with proper
autocomplete and documentation.
burn-book/src/building-blocks/config.md
+64
View File
@@ -0,0 +1,64 @@
# Config
When writing scientific code, you normally have a lot of values that are set, and Deep Learning is
no exception. Python has the possibility to define default parameters for functions, which helps
improve the developer experience. However, this has the downside of potentially breaking your code
when upgrading to a new version, as the default values might change without your knowledge, making
debugging very challenging.
With that in mind, we came up with the Config system. It's a simple Rust derive that you can apply
to your types, allowing you to define default values with ease. Additionally, all configs can be
serialized, reducing potential bugs when upgrading versions and improving reproducibility.
```rust , ignore
use burn::config::Config;
#[derive(Config)]
pub struct MyModuleConfig {
d_model: usize,
d_ff: usize,
#[config(default = 0.1)]
dropout: f64,
}
```
The derive also adds useful `with_` methods for every attribute of your config, similar to a builder
pattern, along with a `save` method.
```rust, ignore
fn main() {
let config = MyModuleConfig::new(512, 2048);
println!("{}", config.d_model); // 512
println!("{}", config.d_ff); // 2048
println!("{}", config.dropout); // 0.1
let config = config.with_dropout(0.2);
println!("{}", config.dropout); // 0.2
config.save("config.json").unwrap();
}
```
## Good practices
By using the config type it is easy to create new module instances. The initialization method should
be implemented on the config type with the device as argument.
```rust, ignore
impl MyModuleConfig {
/// Create a module on the given device.
pub fn init<B: Backend>(&self, device: &B::Device) -> MyModule {
MyModule {
linear: LinearConfig::new(self.d_model, self.d_ff).init(device),
dropout: DropoutConfig::new(self.dropout).init(),
}
}
}
```
Then we could add this line to the above `main`:
```rust, ignore
use burn::backend::Wgpu;
let device = Default::default();
let my_module = config.init::<Wgpu>(&device);
```
burn-book/src/building-blocks/dataset.md
+531
View File
@@ -0,0 +1,531 @@
# Dataset
At its core, a dataset is a collection of data typically related to a specific analysis or
processing task. The data modality can vary depending on the task, but most datasets primarily
consist of images, texts, audio or videos.
This data source represents an integral part of machine learning to successfully train a model.
Thus, it is essential to provide a convenient and performant API to handle your data. Since this
process varies wildly from one problem to another, it is defined as a trait that should be
implemented on your type. The dataset trait is quite similar to the dataset abstract class in
PyTorch:
```rust, ignore
pub trait Dataset<I>: Send + Sync {
fn get(&self, index: usize) -> Option<I>;
fn len(&self) -> usize;
}
```
The dataset trait assumes a fixed-length set of items that can be randomly accessed in constant
time. This is a major difference from datasets that use Apache Arrow underneath to improve streaming
performance. Datasets in Burn don't assume _how_ they are going to be accessed; it's just a
collection of items.
However, you can compose multiple dataset transformations to lazily obtain what you want with zero
pre-processing, so that your training can start instantly!
## Transformation
Transformations in Burn are all lazy and modify one or multiple input datasets. The goal of these
transformations is to provide you with the necessary tools so that you can model complex data
distributions.
| Transformation | Description |
| ------------------ | ------------------------------------------------------------------------------------------------------------------------ |
| `SamplerDataset` | Samples items from a dataset. This is a convenient way to model a dataset as a probability distribution of a fixed size. |
| `SelectionDataset` | Selects a subset of items by index from a dataset. Can be randomly shuffled; can be re-shuffled. |
| `ShuffledDataset` | Shuffles a wrapped dataset; This is a thin wrapper around `SelectionDataset`. |
| `PartialDataset` | Returns a view of the input dataset with a specified range. |
| `MapperDataset` | Computes a transformation lazily on the input dataset. |
| `ComposedDataset` | Composes multiple datasets together to create a larger one without copying any data. |
| `WindowsDataset` | Dataset designed to work with overlapping windows of data extracted from an input dataset. |
Let us look at the basic usages of each dataset transform and how they can be composed together.
These transforms are lazy by default except when specified, reducing the need for unnecessary
intermediate allocations and improving performance. The full documentation of each transform can be
found at the [API reference](https://burn.dev/docs/burn/data/dataset/transform/index.html).
- **SamplerDataset**: This transform can be used to sample items from a dataset with (default) or
without replacement. Transform is initialized with a sampling size which can be bigger or smaller
than the input dataset size. This is particularly useful in cases where we want to checkpoint
larger datasets more often during training and smaller datasets less often as the size of an epoch
is now controlled by the sampling size. Sample usage:
```rust, ignore
type DbPedia = SqliteDataset<DbPediaItem>;
let dataset: DbPedia = HuggingfaceDatasetLoader::new("fancyzhx/dbpedia_14")
.dataset("train").
.unwrap();
let dataset = SamplerDataset<DbPedia, DbPediaItem>::new(dataset, 10000);
```
- **SelectionDataset**: This transform can be used to select a subset of items from a dataset by
index. It can be initialized with a list of indices to select from the input dataset. This is
particularly useful when you want to create a smaller dataset from a larger one, for example, to
create a validation set from a training set.
The `SelectionDataset` can also be initialized with a random seed to shuffle the indices before
selection. This is useful when you want to randomly select a subset of items from the dataset.
Base dataset items may be included more than once in the selection.
```rust, ignore
let explicit = SelectionDataset::from_indices_checked(dataset.clone(), vec![0, 1, 2, 0]);
let shuffled = SelectionDataset::new_shuffled(dataset.clone(), &mut rng);
let shuffled = SelectionDataset::new_shuffled(dataset.clone(), 42);
let mut mutable = SelectionDataset::new_select_all(dataset.clone(), vec![0, 1, 2, 0]);
mutable.shuffle(42);
mutable.shuffle(&mut rng);
```
- **ShuffledDataset**: This transform can be used to shuffle the items of a dataset. Particularly
useful before splitting the raw dataset into train/test splits. Can be initialized with a seed to
ensure reproducibility.
The `ShuffledDataset` is a thin wrapper around the `SelectionDataset`.
```rust, ignore
let dataset = ShuffledDataset<DbPedia, DbPediaItem>::new(dataset, &mut rng);
let dataset = ShuffledDataset<DbPedia, DbPediaItem>::new(dataset, 42);
```
- **PartialDataset**: This transform is useful to return a view of the dataset with specified start
and end indices. Used to create train/val/test splits. In the example below, we show how to chain
ShuffledDataset and PartialDataset to create splits.
```rust, ignore
// define chained dataset type here for brevity
type PartialData = PartialDataset<ShuffledDataset<DbPedia, DbPediaItem>>;
let len = dataset.len();
let split = "train"; // or "val"/"test"
let data_split = match split {
"train" => PartialData::new(dataset, 0, len * 8 / 10), // Get first 80% dataset
"test" => PartialData::new(dataset, len * 8 / 10, len), // Take remaining 20%
_ => panic!("Invalid split type"), // Handle unexpected split types
};
```
- **MapperDataset**: This transform is useful to apply a transformation on each of the items of a
dataset. Particularly useful for normalization of image data when channel means are known.
- **ComposedDataset**: This transform is useful to compose multiple datasets downloaded from
multiple sources (say different HuggingfaceDatasetLoader sources) into a single bigger dataset
which can be sampled from one source.
- **WindowsDataset**: This transform is useful to create overlapping windows of a dataset.
Particularly useful for sequential Time series Data, for example when working with an LSTM.
## Storage
There are multiple dataset storage options available for you to choose from. The choice of the
dataset to use should be based on the dataset's size as well as its intended purpose.
| Storage | Description |
| ------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------- |
| `InMemDataset` | In-memory dataset that uses a vector to store items. Well-suited for smaller datasets. |
| `SqliteDataset` | Dataset that uses [SQLite](https://www.sqlite.org/) to index items that can be saved in a simple SQL database file. Well-suited for larger datasets. |
| `DataframeDataset` | Dataset that uses [Polars](https://www.pola.rs/) dataframe to store and manage data. Well-suited for efficient data manipulation and analysis. |
## Sources
For now, there are only a couple of dataset sources available with Burn, but more to come!
### Hugging Face
You can easily import any Hugging Face dataset with Burn. We use SQLite as the storage to avoid
downloading the model each time or starting a Python process. You need to know the format of each
item in the dataset beforehand. Here's an example with the
[dbpedia dataset](https://huggingface.co/datasets/fancyzhx/dbpedia_14).
```rust, ignore
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct DbPediaItem {
pub title: String,
pub content: String,
pub label: usize,
}
fn main() {
let dataset: SqliteDataset<DbPediaItem> = HuggingfaceDatasetLoader::new("fancyzhx/dbpedia_14")
.dataset("train") // The training split.
.unwrap();
}
```
We see that items must derive `serde::Serialize`, `serde::Deserialize`, `Clone`, and `Debug`, but
those are the only requirements.
<div class="warning">
The `HuggingfaceDatasetLoader` relies on the
[`datasets` library by HuggingFace](https://huggingface.co/docs/datasets/index) to download
datasets. This is a Python library, so you must have an existing Python installation to use this
loader.
</div>
### Images
`ImageFolderDataset` is a generic vision dataset used to load images from disk. It is currently
available for multi-class and multi-label classification tasks as well as semantic segmentation and
object detection tasks.
```rust, ignore
// Create an image classification dataset from the root folder,
// where images for each class are stored in their respective folder.
//
// For example:
// root/dog/dog1.png
// root/dog/dog2.png
// ...
// root/cat/cat1.png
let dataset = ImageFolderDataset::new_classification("path/to/dataset/root").unwrap();
```
```rust, ignore
// Create a multi-label image classification dataset from a list of items,
// where each item is a tuple `(image path, labels)`, and a list of classes
// in the dataset.
//
// For example:
let items = vec![
("root/dog/dog1.png", vec!["animal".to_string(), "dog".to_string()]),
("root/cat/cat1.png", vec!["animal".to_string(), "cat".to_string()]),
];
let dataset = ImageFolderDataset::new_multilabel_classification_with_items(
items,
&["animal", "cat", "dog"],
)
.unwrap();
```
```rust, ignore
// Create a segmentation mask dataset from a list of items, where each
// item is a tuple `(image path, mask path)` and a list of classes
// corresponding to the integer values in the mask.
let items = vec![
(
"path/to/images/image0.png",
"path/to/annotations/mask0.png",
),
(
"path/to/images/image1.png",
"path/to/annotations/mask1.png",
),
(
"path/to/images/image2.png",
"path/to/annotations/mask2.png",
),
];
let dataset = ImageFolderDataset::new_segmentation_with_items(
items,
&[
"cat", // 0
"dog", // 1
"background", // 2
],
)
.unwrap();
```
```rust, ignore
// Create an object detection dataset from a COCO dataset. Currently only
// the import of object detection data (bounding boxes) is supported.
//
// COCO offers separate annotation and image archives for training and
// validation, paths to the unpacked files need to be passed as parameters:
let dataset = ImageFolderDataset::new_coco_detection(
"/path/to/coco/instances_train2017.json",
"/path/to/coco/images/train2017"
)
.unwrap();
```
### Comma-Separated Values (CSV)
Loading records from a simple CSV file in-memory is simple with the `InMemDataset`:
```rust, ignore
// Build dataset from csv with tab ('\t') delimiter.
// The reader can be configured for your particular file.
let mut rdr = csv::ReaderBuilder::new();
let rdr = rdr.delimiter(b'\t');
let dataset = InMemDataset::from_csv("path/to/csv", rdr).unwrap();
```
Note that this requires the `csv` crate.
**What about streaming datasets?**
There is no streaming dataset API with Burn, and this is by design! The learner struct will iterate
multiple times over the dataset and only checkpoint when done. You can consider the length of the
dataset as the number of iterations before performing checkpointing and running the validation.
There is nothing stopping you from returning different items even when called with the same `index`
multiple times.
## How Is The Dataset Used?
During training, the dataset is used to access the data samples and, for most use cases in
supervised learning, their corresponding ground-truth labels. Remember that the `Dataset` trait
implementation is responsible to retrieve the data from its source, usually some sort of data
storage. At this point, the dataset could be naively iterated over to provide the model a single
sample to process at a time, but this is not very efficient.
Instead, we collect multiple samples that the model can process as a _batch_ to fully leverage
modern hardware (e.g., GPUs - which have impressive parallel processing capabilities). Since each
data sample in the dataset can be collected independently, the data loading is typically done in
parallel to further speed things up. In this case, we parallelize the data loading using a
multi-threaded `BatchDataLoader` to obtain a sequence of items from the `Dataset` implementation.
Finally, the sequence of items is combined into a batched tensor that can be used as input to a
model with the `Batcher` trait implementation. Other tensor operations can be performed during this
step to prepare the batch data, as is done [in the basic workflow guide](../basic-workflow/data.md).
The process is illustrated in the figure below for the MNIST dataset.
<img title="Burn Data Loading Pipeline" alt="Burn Data Loading Pipeline" src="./dataset.png">
Although we have conveniently implemented the
[`MnistDataset`](https://github.com/tracel-ai/burn/blob/main/crates/burn-dataset/src/vision/mnist.rs)
used in the guide, we'll go over its implementation to demonstrate how the `Dataset` and `Batcher`
traits are used.
The [MNIST dataset](http://yann.lecun.com/exdb/mnist/) of handwritten digits has a training set of
60,000 examples and a test set of 10,000 examples. A single item in the dataset is represented by a
\\(28 \times 28\\) pixels black-and-white image (stored as raw bytes) with its corresponding label
(a digit between \\(0\\) and \\(9\\)). This is defined by the `MnistItemRaw` struct.
```rust, ignore
# #[derive(Deserialize, Debug, Clone)]
struct MnistItemRaw {
pub image_bytes: Vec<u8>,
pub label: u8,
}
```
With single-channel images of such low resolution, the entire training and test sets can be loaded
in memory at once. Therefore, we leverage the already existing `InMemDataset` to retrieve the raw
images and labels data. At this point, the image data is still just a bunch of bytes, but we want to
retrieve the _structured_ image data in its intended form. For that, we can define a `MapperDataset`
that transforms the raw image bytes to a 2D array image (which we convert to float while we're at
it).
```rust, ignore
const WIDTH: usize = 28;
const HEIGHT: usize = 28;
# /// MNIST item.
# #[derive(Deserialize, Serialize, Debug, Clone)]
pub struct MnistItem {
/// Image as a 2D array of floats.
pub image: [[f32; WIDTH]; HEIGHT],
/// Label of the image.
pub label: u8,
}
struct BytesToImage;
impl Mapper<MnistItemRaw, MnistItem> for BytesToImage {
/// Convert a raw MNIST item (image bytes) to a MNIST item (2D array image).
fn map(&self, item: &MnistItemRaw) -> MnistItem {
// Ensure the image dimensions are correct.
debug_assert_eq!(item.image_bytes.len(), WIDTH * HEIGHT);
// Convert the image to a 2D array of floats.
let mut image_array = [[0f32; WIDTH]; HEIGHT];
for (i, pixel) in item.image_bytes.iter().enumerate() {
let x = i % WIDTH;
let y = i / HEIGHT;
image_array[y][x] = *pixel as f32;
}
MnistItem {
image: image_array,
label: item.label,
}
}
}
type MappedDataset = MapperDataset<InMemDataset<MnistItemRaw>, BytesToImage, MnistItemRaw>;
# /// The MNIST dataset consists of 70,000 28x28 black-and-white images in 10 classes (one for each digits), with 7,000
# /// images per class. There are 60,000 training images and 10,000 test images.
# ///
# /// The data is downloaded from the web from the [CVDF mirror](https://github.com/cvdfoundation/mnist).
pub struct MnistDataset {
dataset: MappedDataset,
}
```
To construct the `MnistDataset`, the data source must be parsed into the expected `MappedDataset`
type. Since both the train and test sets use the same file format, we can separate the functionality
to load the `train()` and `test()` dataset.
```rust, ignore
impl MnistDataset {
/// Creates a new train dataset.
pub fn train() -> Self {
Self::new("train")
}
/// Creates a new test dataset.
pub fn test() -> Self {
Self::new("test")
}
fn new(split: &str) -> Self {
// Download dataset
let root = MnistDataset::download(split);
// Parse data as vector of images bytes and vector of labels
let images: Vec<Vec<u8>> = MnistDataset::read_images(&root, split);
let labels: Vec<u8> = MnistDataset::read_labels(&root, split);
// Collect as vector of MnistItemRaw
let items: Vec<_> = images
.into_iter()
.zip(labels)
.map(|(image_bytes, label)| MnistItemRaw { image_bytes, label })
.collect();
// Create the MapperDataset for InMemDataset<MnistItemRaw> to transform
// items (MnistItemRaw -> MnistItem)
let dataset = InMemDataset::new(items);
let dataset = MapperDataset::new(dataset, BytesToImage);
Self { dataset }
}
# /// Download the MNIST dataset files from the web.
# /// Panics if the download cannot be completed or the content of the file cannot be written to disk.
# fn download(split: &str) -> PathBuf {
# // Dataset files are stored in the burn-dataset cache directory
# let cache_dir = dirs::cache_dir()
# .expect("Could not get cache directory")
# .join("burn-dataset");
# let split_dir = cache_dir.join("mnist").join(split);
#
# if !split_dir.exists() {
# create_dir_all(&split_dir).expect("Failed to create base directory");
# }
#
# // Download split files
# match split {
# "train" => {
# MnistDataset::download_file(TRAIN_IMAGES, &split_dir);
# MnistDataset::download_file(TRAIN_LABELS, &split_dir);
# }
# "test" => {
# MnistDataset::download_file(TEST_IMAGES, &split_dir);
# MnistDataset::download_file(TEST_LABELS, &split_dir);
# }
# _ => panic!("Invalid split specified {}", split),
# };
#
# split_dir
# }
#
# /// Download a file from the MNIST dataset URL to the destination directory.
# /// File download progress is reported with the help of a [progress bar](indicatif).
# fn download_file<P: AsRef<Path>>(name: &str, dest_dir: &P) -> PathBuf {
# // Output file name
# let file_name = dest_dir.as_ref().join(name);
#
# if !file_name.exists() {
# // Download gzip file
# let bytes = download_file_as_bytes(&format!("{URL}{name}.gz"), name);
#
# // Create file to write the downloaded content to
# let mut output_file = File::create(&file_name).unwrap();
#
# // Decode gzip file content and write to disk
# let mut gz_buffer = GzDecoder::new(&bytes[..]);
# std::io::copy(&mut gz_buffer, &mut output_file).unwrap();
# }
#
# file_name
# }
#
# /// Read images at the provided path for the specified split.
# /// Each image is a vector of bytes.
# fn read_images<P: AsRef<Path>>(root: &P, split: &str) -> Vec<Vec<u8>> {
# let file_name = if split == "train" {
# TRAIN_IMAGES
# } else {
# TEST_IMAGES
# };
# let file_name = root.as_ref().join(file_name);
#
# // Read number of images from 16-byte header metadata
# let mut f = File::open(file_name).unwrap();
# let mut buf = [0u8; 4];
# let _ = f.seek(SeekFrom::Start(4)).unwrap();
# f.read_exact(&mut buf)
# .expect("Should be able to read image file header");
# let size = u32::from_be_bytes(buf);
#
# let mut buf_images: Vec<u8> = vec![0u8; WIDTH * HEIGHT * (size as usize)];
# let _ = f.seek(SeekFrom::Start(16)).unwrap();
# f.read_exact(&mut buf_images)
# .expect("Should be able to read image file header");
#
# buf_images
# .chunks(WIDTH * HEIGHT)
# .map(|chunk| chunk.to_vec())
# .collect()
# }
#
# /// Read labels at the provided path for the specified split.
# fn read_labels<P: AsRef<Path>>(root: &P, split: &str) -> Vec<u8> {
# let file_name = if split == "train" {
# TRAIN_LABELS
# } else {
# TEST_LABELS
# };
# let file_name = root.as_ref().join(file_name);
#
# // Read number of labels from 8-byte header metadata
# let mut f = File::open(file_name).unwrap();
# let mut buf = [0u8; 4];
# let _ = f.seek(SeekFrom::Start(4)).unwrap();
# f.read_exact(&mut buf)
# .expect("Should be able to read label file header");
# let size = u32::from_be_bytes(buf);
#
# let mut buf_labels: Vec<u8> = vec![0u8; size as usize];
# let _ = f.seek(SeekFrom::Start(8)).unwrap();
# f.read_exact(&mut buf_labels)
# .expect("Should be able to read labels from file");
#
# buf_labels
# }
}
```
Since the `MnistDataset` simply wraps a `MapperDataset` instance with `InMemDataset`, we can easily
implement the `Dataset` trait.
```rust, ignore
impl Dataset<MnistItem> for MnistDataset {
fn get(&self, index: usize) -> Option<MnistItem> {
self.dataset.get(index)
}
fn len(&self) -> usize {
self.dataset.len()
}
}
```
The only thing missing now is the `Batcher`, which we already went over
[in the basic workflow guide](../basic-workflow/data.md). The `Batcher` takes a list of `MnistItem`
retrieved by the dataloader as input and returns a batch of images as a 3D tensor along with their
targets.
burn-book/src/building-blocks/dataset.png
Executable
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 99 KiB

burn-book/src/building-blocks/learner.md
+112
View File
@@ -0,0 +1,112 @@
# Learner
The [burn-train](https://github.com/tracel-ai/burn/tree/main/crates/burn-train) crate encapsulates
multiple utilities for training deep learning models. The goal of the crate is to provide users with
a well-crafted and flexible training loop, so that projects do not have to write such components
from the ground up. Most of the interactions with `burn-train` will be with the `SupervisedTraining`
struct, briefly presented in the previous [training section](../basic-workflow/training.md). This
struct enables you to configure the training loop, offering support for registering metrics,
enabling logging, checkpointing states, using multiple devices, and so on.
There are still some assumptions in the current provided APIs, which may make them inappropriate for
your learning requirements. Indeed, they assume your model will learn from a training dataset and be
validated against another dataset. This is the most common paradigm, allowing users to do both
supervised and unsupervised learning as well as fine-tuning. However, for more complex requirements,
creating a [custom training loop](../custom-training-loop.md) might be what you need.
## Usage
The `SupervisedLearning` struct must be created with the training and validation dataloaders. It provides you with numerous options when it comes to configurations.
| Configuration | Description |
| ---------------------- | ------------------------------------------------------------------------------ |
| Training Metric | Register a training metric |
| Validation Metric | Register a validation metric |
| Training Metric Plot | Register a training metric with plotting (requires the metric to be numeric) |
| Validation Metric Plot | Register a validation metric with plotting (requires the metric to be numeric) |
| Metric Logger | Configure the metric loggers (default is saving them to files) |
| Renderer | Configure how to render metrics (default is CLI) |
| Grad Accumulation | Configure the number of steps before applying gradients |
| File Checkpointer | Configure how the model, optimizer and scheduler states are saved |
| Num Epochs | Set the number of epochs |
| Devices | Set the devices to be used |
| Checkpoint | Restart training from a checkpoint |
| Application logging | Configure the application logging installer (default is writing to `experiment.log`) |
| Training Strategy | Use a custom training strategy, allowing you to use your own training loop with all the capabilities of the `SupervisedTraining` struct |
When the training is configured to your liking, you can then move forward to running the training. The
`launch` method requires a learner object providing: the model, the optimizer and the learning rate scheduler. Note
that the latter can be a simple float if you want it to be constant during training.
The `launch` method will start the training and return the trained model once finished.
Again, please refer to the [training section](../basic-workflow/training.md) for a relevant code
snippet.
## Multiple optimizers
It's common practice to set different learning rates, optimizer parameters, or use different optimizers entirely, for different parts
of a model. You can leverage Burn's `ParamGroup`s to mix and match optimizers and learning rate schedulers easily!
```rust,ignore
let lr_scheduler_base = ComposedLrSchedulerConfig::new()
.cosine(CosineAnnealingLrSchedulerConfig::new(1.0, 2000))
.linear(LinearLrSchedulerConfig::new(1e-8, 1.0, 2000))
.linear(LinearLrSchedulerConfig::new(1e-2, 1e-6, 10000));
let lr_scheduler = lr_scheduler_base.init().unwrap().with_group(
ParamGroup::from_predicate("conv"),
LinearLrSchedulerConfig::new(1e-6, 1e-3, 14000)
.build()
.unwrap(),
);
let optimizer_base = AdamWConfig::new()
.with_cautious_weight_decay(true)
.with_weight_decay(5e-5);
let optim = optimizer_base.init().with_group(
ParamGroup::from_predicate("conv"),
SgdConfig::new().build(),
None,
);
let result = training.launch(Learner::new(
model,
optim,
lr_scheduler,
));
```
## Artifacts
When creating a `SupervisedTraining` instance, all the collected data will be saved under the directory provided as
the argument to the `new` method. Here is an example of the data layout for a model checkpointed to
the burnpack format, with the accuracy and loss metrics registered:
```
├── experiment.log
├── checkpoint
│   ├── model-1.bpk
│   ├── optim-1.bpk
│   └── scheduler-1.bpk
│   ├── model-2.bpk
│   ├── optim-2.bpk
│   └── scheduler-2.bpk
├── train
│   ├── epoch-1
│   │   ├── Accuracy.log
│   │   └── Loss.log
│   └── epoch-2
│   ├── Accuracy.log
│   └── Loss.log
└── valid
├── epoch-1
│   ├── Accuracy.log
│   └── Loss.log
└── epoch-2
├── Accuracy.log
└── Loss.log
```
You can choose to save or synchronize that local directory with a remote file system, if desired.
The file checkpointer is capable of automatically deleting old checkpoints according to a specified
configuration.
burn-book/src/building-blocks/metric.md
+226
View File
@@ -0,0 +1,226 @@
# Metric
When working with the learner, you have the option to record metrics that will be monitored
throughout the training process. We currently offer a restricted range of metrics.
| Metric | Description |
| ------------------- | ------------------------------------------------------------------------------------------- |
| Accuracy | Calculate the accuracy in percentage |
| TopKAccuracy | Calculate the top-k accuracy in percentage |
| Precision | Calculate precision in percentage |
| Recall | Calculate recall in percentage |
| FBetaScore | Calculate F<sub>β </sub>score in percentage |
| AUROC | Calculate the area under curve of ROC in percentage |
| AUC-PR | Calculate the area under the precision-recall curve (average precision) in percentage |
| Loss | Output the loss used for the backward pass |
| CharErrorRate (CER) | Calculate Character Error Rate in percentage |
| WordErrorRate (WER) | Calculate Word Error Rate in percentage |
| HammingScore | Calculate hamming score (also known as multi-label or label-based accuracy) in percentage |
| Perplexity | Calculate perplexity which is a measure of how well a probability model predicts samples |
| IterationSpeed | Tracks the training iteration speed, measuring how many iterations are completed per second |
| CPU Temperature | Fetch the temperature of CPUs |
| CPU Usage | Fetch the CPU utilization |
| CPU Memory Usage | Fetch the CPU RAM usage |
| Learning Rate | Fetch the current learning rate for each optimizer step |
| CUDA | Fetch general CUDA metrics such as utilization |
| Vision Metric | Description |
| ------------- | ---------------------------------------------------------------------------------------------------- |
| A-FINE | Computes the Adaptive Fidelity-Naturalness Evaluator (A-FINE) full-reference perceptual quality metric built on CLIP ViT-B/32 features |
| Dice | Computes the Dice-Sorenson coefficient (DSC) for evaluating overlap between binary masks |
| DISTS | Computes the Deep Image Structure and Texture Similarity (DISTS) metric for image quality assessment |
| FID | Computes the Frechet Inception Distance (FID) for evaluating generative model quality |
| LPIPS | Computes the Learned Perceptual Image Patch Similarity (LPIPS) for image quality assessment |
| MS-SSIM | Computes the Multi-scale Structural Similarity index measure (MS-SSIM) for image quality assessment |
| PSNR | Computes the Peak Signal-to-Noise Ratio (PSNR) for image quality assessment |
| SSIM | Computes the Structural Similarity index measure (SSIM) for image quality assessment |
## Using Metrics with the Learner
In order to use a metric, the output of your training step must implement the `Adaptor` trait from
`burn-train::metric` for each metric's corresponding input type. The `Adaptor` trait simply converts
your output struct into the input type the metric expects.
Burn provides four built-in output structs that cover common tasks. Each one already implements
`Adaptor` for a set of metrics, so in many cases you can use them directly without writing any
adaptor code yourself.
- `ClassificationOutput<B>`:
- Use case: Single-label classification
- Fields: `loss: Tensor<B, 1>`, `output: Tensor<B, 2>`, `targets: Tensor<B, 1, Int>`
- Adapted metrics: Accuracy, TopKAccuracy, Perplexity, Precision\*, Recall\*, FBetaScore\*, AUROC\*, AUC-PR\*, Loss
- `MultiLabelClassificationOutput<B>`:
- Use case: Multi-label classification
- Fields: `loss: Tensor<B, 1>`, `output: Tensor<B, 2>`, `targets: Tensor<B, 2, Int>`
- Adapted metrics: HammingScore, Precision\*, Recall\*, FBetaScore\*, AUROC\*, AUC-PR\*, Loss
- `RegressionOutput<B>`:
- Use case: Regression tasks
- Fields: `loss: Tensor<B, 1>`, `output: Tensor<B, 2>`, `targets: Tensor<B, 2>`
- Adapted metrics: Loss
- `SequenceOutput<B>`:
- Use case: Sequence prediction
- Fields: `loss: Tensor<B, 1>`, `logits: Tensor<B, 3>`, `predictions: Option<Tensor<B, 2, Int>>`, `targets: Tensor<B, 2, Int>`
- Adapted metrics: Accuracy, TopKAccuracy, Perplexity, CER, WER, Loss
\* Precision, Recall, FBetaScore, AUROC, and AUC-PR all use `ConfusionStatsInput` as their input type so these
metrics are automatically (implicitly) adapted since `ConfusionStatsInput` is adapted.
If your metric isn't already adapted for the appropriate output struct, you can implement `Adaptor` yourself.
For example, here is how `ClassificationOutput` adapts to `AccuracyInput`:
```rust,ignore
impl<B: Backend> Adaptor<AccuracyInput<B>> for ClassificationOutput<B> {
fn adapt(&self) -> AccuracyInput<B> {
AccuracyInput::new(self.output.clone(), self.targets.clone())
}
}
```
If your task type is not covered by the built-in output structs, you can create an output struct for your data
and then adapt your metric for the output struct:
```rust,ignore
#[derive(new)]
pub struct ClassificationOutput<B: Backend> {
/// The loss.
pub loss: Tensor<B, 1>,
/// The output.
pub output: Tensor<B, 2>,
/// The targets.
pub targets: Tensor<B, 1, Int>,
}
impl<B: Backend> Adaptor<AccuracyInput<B>> for ClassificationOutput<B> {
fn adapt(&self) -> AccuracyInput<B> {
AccuracyInput::new(self.output.clone(), self.targets.clone())
}
}
```
You can also open an issue on the [GitHub repository](https://github.com/tracel-ai/burn) when your task type is
not covered by the built-in output structs. However, since creating an output struct for your data is simple,
it is recommended to try creating your own output struct first.
# Custom Metric
Generating your own custom metrics is done by implementing the `Metric` trait.
```rust , ignore
/// Metric trait.
///
/// # Notes
///
/// Implementations should define their own input type only used by the metric.
/// This is important since some conflict may happen when the model output is adapted for each
/// metric's input type.
pub trait Metric: Send + Sync + Clone {
/// The input type of the metric.
type Input;
/// The parameterized name of the metric.
///
/// This should be unique, so avoid using short generic names, prefer using the long name.
///
/// For a metric that can exist at different parameters (e.g., top-k accuracy for different
/// values of k), the name should be unique for each instance.
fn name(&self) -> MetricName;
/// Update the metric state and returns the current metric entry.
fn update(&mut self, item: &Self::Input, metadata: &MetricMetadata) -> SerializedEntry;
/// Clear the metric state.
fn clear(&mut self);
}
```
As an example, let's see how the loss metric is implemented.
```rust, ignore
/// The loss metric.
#[derive(Clone)]
pub struct LossMetric<B: Backend> {
name: Arc<String>,
state: NumericMetricState,
_b: B,
}
/// The [loss metric](LossMetric) input type.
#[derive(new)]
pub struct LossInput<B: Backend> {
tensor: Tensor<B, 1>,
}
impl<B: Backend> Default for LossMetric<B> {
fn default() -> Self {
Self::new()
}
}
impl<B: Backend> LossMetric<B> {
/// Create the metric.
pub fn new() -> Self {
Self {
name: Arc::new("Loss".to_string()),
state: NumericMetricState::default(),
_b: Default::default(),
}
}
}
impl<B: Backend> Metric for LossMetric<B> {
type Input = LossInput<B>;
fn update(&mut self, loss: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry {
let [batch_size] = loss.tensor.dims();
let loss = loss
.tensor
.clone()
.mean()
.into_data()
.iter::<f64>()
.next()
.unwrap();
self.state.update(
loss,
batch_size,
FormatOptions::new(self.name()).precision(2),
)
}
fn clear(&mut self) {
self.state.reset()
}
fn name(&self) -> MetricName {
self.name.clone()
}
fn attributes(&self) -> MetricAttributes {
NumericAttributes {
unit: None,
higher_is_better: false,
}
.into()
}
}
```
When the metric you are implementing is numeric in nature, you may want to also implement the
`Numeric` trait. This will allow your metric to be plotted.
```rust, ignore
impl<B: Backend> Numeric for LossMetric<B> {
fn value(&self) -> NumericEntry {
self.state.current_value()
}
fn running_value(&self) -> NumericEntry {
self.state.running_value()
}
}
```
burn-book/src/building-blocks/module.md
+343
View File
@@ -0,0 +1,343 @@
# Module
The `Module` derive allows you to create your own neural network modules, similar to PyTorch. The
derive function only generates the necessary methods to essentially act as a parameter container for
your type, it makes no assumptions about how the forward pass is declared.
```rust, ignore
use burn::module::Module;
use burn::tensor::backend::Backend;
#[derive(Module, Debug)]
pub struct PositionWiseFeedForward<B: Backend> {
linear_inner: Linear<B>,
linear_outer: Linear<B>,
dropout: Dropout,
gelu: Gelu,
}
impl<B: Backend> PositionWiseFeedForward<B> {
/// Normal method added to a struct.
pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
let x = self.linear_inner.forward(input);
let x = self.gelu.forward(x);
let x = self.dropout.forward(x);
self.linear_outer.forward(x)
}
}
```
Note that all fields declared in the struct must also implement the `Module` trait.
## Tensor
If you want to create your own module that contains tensors, and not just other modules defined with
the `Module` derive, you need to be careful to achieve the behavior you want.
- `Param<Tensor<B, D>>`: If you want the tensor to be included as a parameter of your modules, you
need to wrap the tensor in a `Param` struct. This will create an ID that will be used to identify
this parameter. This is essential when performing module optimization and when saving states such
as optimizer and module checkpoints. Note that a module's record only contains parameters.
- `Param<Tensor<B, D>>.set_require_grad(false)`: If you want the tensor to be included as a
parameter of your modules, and therefore saved with the module's weights, but you don't want it to
be updated by the optimizer.
- `Tensor<B, D>`: If you want the tensor to act as a constant that can be recreated when
instantiating a module. This can be useful when generating sinusoidal embeddings, for example.
## Methods
These methods are available for all modules.
| Burn API | PyTorch Equivalent |
| --------------------------------------- | ---------------------------------------- |
| `module.devices()` | N/A |
| `module.fork(device)` | Similar to `module.to(device).detach()` |
| `module.to_device(device)` | `module.to(device)` |
| `module.no_grad()` | `module.require_grad_(False)` |
| `module.num_params()` | N/A |
| `module.visit(visitor)` | N/A |
| `module.map(mapper)` | N/A |
| `module.into_record()` | Similar to `state_dict` |
| `module.load_record(record)` | Similar to `load_state_dict(state_dict)` |
| `module.into_record().save(file_path)` | Similar to `torch.save(state_dict, ...)` |
| `ModuleRecord::load(file_path)` | Similar to `torch.load(...)` |
Similar to the backend trait, there is also the `AutodiffModule` trait to signify a module with
autodiff support.
| Burn API | PyTorch Equivalent |
| ---------------- | ------------------ |
| `module.valid()` | `module.eval()` |
## Visitor & Mapper
As mentioned earlier, modules primarily function as parameter containers. Therefore, we naturally
offer several ways to perform functions on each parameter. This is distinct from PyTorch, where
extending module functionalities is not as straightforward.
The `map` and `visitor` methods are quite similar but serve different purposes. Mapping is used for
potentially mutable operations where each parameter of a module can be updated to a new value. In
Burn, optimizers are essentially just sophisticated module mappers. Visitors, on the other hand, are
used when you don't intend to modify the module but need to retrieve specific information from it,
such as the number of parameters or a list of devices in use.
You can implement your own mapper or visitor by implementing these simple traits:
```rust, ignore
/// Module visitor trait.
pub trait ModuleVisitor<B: Backend> {
/// Visit a float tensor in the module.
fn visit_float<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D>);
/// Visit an int tensor in the module.
fn visit_int<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D, Int>);
/// Visit a bool tensor in the module.
fn visit_bool<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D, Bool>);
}
/// Module mapper trait.
pub trait ModuleMapper<B: Backend> {
/// Map a float tensor in the module.
fn map_float<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D>) -> Tensor<B, D>;
/// Map an int tensor in the module.
fn map_int<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D, Int>) -> Tensor<B, D, Int>;
/// Map a bool tensor in the module.
fn map_bool<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D, Bool>) -> Tensor<B, D, Bool>;
}
```
Note that the trait doesn't require all methods to be implemented as they are already defined to
perform no operation. If you're only interested in float tensors (like the majority of use cases),
then you can simply implement `map_float` or `visit_float`.
For example, the `ModuleMapper` trait could be implemented to clamp all parameters into the range
`[min, max]`.
```rust, ignore
/// Clamp parameters into the range `[min, max]`.
pub struct Clamp {
/// Lower-bound of the range.
pub min: f32,
/// Upper-bound of the range.
pub max: f32,
}
// Clamp all floating-point parameter tensors between `[min, max]`.
impl<B: Backend> ModuleMapper<B> for Clamp {
fn map_float<const D: usize>(
&mut self,
_id: burn::module::ParamId,
tensor: burn::prelude::Tensor<B, D>,
) -> burn::prelude::Tensor<B, D> {
tensor.clamp(self.min, self.max)
}
}
// Clamp module mapper into the range `[-0.5, 0.5]`
let mut clamp = Clamp {
min: -0.5,
max: 0.5,
};
let model = model.map(&mut clamp);
```
If you want to use this during training to constrain your model parameters, make sure that the
parameter tensors are still tracked for autodiff. This can be done with a simple adjustment to the
implementation.
```rust, ignore
impl<B: AutodiffBackend> ModuleMapper<B> for Clamp {
fn map_float<const D: usize>(
&mut self,
_id: burn::module::ParamId,
tensor: burn::prelude::Tensor<B, D>,
) -> burn::prelude::Tensor<B, D> {
let is_require_grad = tensor.is_require_grad();
let mut tensor = Tensor::from_inner(tensor.inner().clamp(self.min, self.max));
if is_require_grad {
tensor = tensor.require_grad();
}
tensor
}
}
```
## Module Display
Burn provides a simple way to display the structure of a module and its configuration at a glance.
You can print the module to see its structure, which is useful for debugging and tracking changes
across different versions of a module. (See the print output of the
[Basic Workflow Model](../basic-workflow/model.md) example.)
To customize the display of a module, you can implement the `ModuleDisplay` trait for your module.
This will change the default display settings for the module and its children. Note that
`ModuleDisplay` is automatically implemented for all modules, but you can override it to customize
the display by annotating the module with `#[module(custom_display)]`.
```rust
#[derive(Module, Debug)]
#[module(custom_display)]
pub struct PositionWiseFeedForward<B: Backend> {
linear_inner: Linear<B>,
linear_outer: Linear<B>,
dropout: Dropout,
gelu: Gelu,
}
impl<B: Backend> ModuleDisplay for PositionWiseFeedForward<B> {
/// Custom settings for the display of the module.
/// If `None` is returned, the default settings will be used.
fn custom_settings(&self) -> Option<burn::module::DisplaySettings> {
DisplaySettings::new()
// Will show all attributes (default is false)
.with_show_all_attributes(false)
// Will show each attribute on a new line (default is true)
.with_new_line_after_attribute(true)
// Will show the number of parameters (default is true)
.with_show_num_parameters(true)
// Will indent by 2 spaces (default is 2)
.with_indentation_size(2)
// Will show the parameter ID (default is false)
.with_show_param_id(false)
// Convenience method to wrap settings in Some()
.optional()
}
/// Custom content to be displayed.
/// If `None` is returned, the default content will be used
/// (all attributes of the module)
fn custom_content(&self, content: Content) -> Option<Content> {
content
.add("linear_inner", &self.linear_inner)
.add("linear_outer", &self.linear_outer)
.add("anything", "anything_else")
.optional()
}
}
```
## Built-in Modules
Burn comes with built-in modules that you can use to build your own modules.
### General
| Burn API | PyTorch Equivalent |
| ----------------- | --------------------------------------------- |
| `BatchNorm` | `nn.BatchNorm1d`, `nn.BatchNorm2d` etc. |
| `Celu` | `nn.CELU` |
| `Dropout` | `nn.Dropout` |
| `Elu` | `nn.ELU` |
| `Embedding` | `nn.Embedding` |
| `GaussianNoise` | _No direct equivalent_ |
| `Gelu` | `nn.Gelu` |
| `Glu` | `nn.Glu` |
| `GroupNorm` | `nn.GroupNorm` |
| `HardShrink` | `nn.Hardshrink` |
| `HardSigmoid` | `nn.Hardsigmoid` |
| `HardSwish` | `nn.Hardswish` |
| `InstanceNorm` | `nn.InstanceNorm1d`, `nn.InstanceNorm2d` etc. |
| `LayerNorm` | `nn.LayerNorm` |
| `LocalResponseNorm` | `nn.LocalResponseNorm` |
| `LeakyRelu` | `nn.LeakyReLU` |
| `Linear` | `nn.Linear` |
| `Prelu` | `nn.PReLu` |
| `Relu` | `nn.ReLU` |
| `Selu` | `nn.SELU` |
| `Sigmoid` | `nn.Sigmoid` |
| `Softplus` | `nn.Softplus` |
| `SoftShrink` | `nn.Softshrink` |
| `Softsign` | `nn.Softsign` |
| `Shrink` | _No direct equivalent_ |
| `RmsNorm` | _No direct equivalent_ |
| `SwiGlu` | _No direct equivalent_ |
| `Tanh` | `nn.Tanh` |
| `ThresholdedRelu` | _No direct equivalent_ |
### Convolutions
| Burn API | PyTorch Equivalent |
| ----------------- | ------------------------------ |
| `Conv1d` | `nn.Conv1d` |
| `Conv2d` | `nn.Conv2d` |
| `Conv3d` | `nn.Conv3d` |
| `ConvTranspose1d` | `nn.ConvTranspose1d` |
| `ConvTranspose2d` | `nn.ConvTranspose2d` |
| `ConvTranspose3d` | `nn.ConvTranspose3d` |
| `DeformConv2d` | `torchvision.ops.DeformConv2d` |
### Pooling
| Burn API | PyTorch Equivalent |
| ------------------- | ---------------------- |
| `AdaptiveAvgPool1d` | `nn.AdaptiveAvgPool1d` |
| `AdaptiveAvgPool2d` | `nn.AdaptiveAvgPool2d` |
| `AvgPool1d` | `nn.AvgPool1d` |
| `AvgPool2d` | `nn.AvgPool2d` |
| `MaxPool1d` | `nn.MaxPool1d` |
| `MaxPool2d` | `nn.MaxPool2d` |
### Interpolation
| Burn API | PyTorch Equivalent |
| --------------- | ------------------ |
| `Interpolate1d` | `nn.Upsample` |
| `Interpolate2d` | `nn.Upsample` |
Interpolation modules resize tensors using one of the available `InterpolateMode` options:
| Mode | Description |
| --------- | -------------------------------------------------------- |
| `Nearest` | Nearest-neighbor interpolation |
| `Linear` | Linear interpolation (bilinear for 2D) |
| `Cubic` | Cubic interpolation (bicubic for 2D) |
| `Lanczos` | Lanczos3 resampling (6-tap sinc-based filter, a=3) |
Configuration is done via `Interpolate1dConfig` / `Interpolate2dConfig` with these options:
| Option | Type | Default | Description |
| --------------- |------------------------------------------| --------- | -------------------------------------------------------- |
| `output_size` | `Option<usize>` / `Option<[usize; 2]>` | `None` | Target output size (takes precedence over scale_factor) |
| `scale_factor` | `Option<f32>` / `Option<[f32; 2]>` | `None` | Scale factor for resizing |
| `mode` | `InterpolateMode` | `Nearest` | Interpolation algorithm |
| `align_corners` | `bool` | `true` | Align input/output corner pixels |
### RNNs
| Burn API | PyTorch Equivalent |
| ---------------- | ---------------------- |
| `Gru`/`BiGru` | `nn.GRU` |
| `Lstm`/`BiLstm` | `nn.LSTM` |
| `GateController` | _No direct equivalent_ |
### Transformer
| Burn API | PyTorch Equivalent |
| -------------------- | ----------------------- |
| `MultiHeadAttention` | `nn.MultiheadAttention` |
| `TransformerDecoder` | `nn.TransformerDecoder` |
| `TransformerEncoder` | `nn.TransformerEncoder` |
| `PositionalEncoding` | _No direct equivalent_ |
| `RotaryEncoding` | _No direct equivalent_ |
### Loss
| Burn API | PyTorch Equivalent |
| ------------------------ | ------------------------ |
| `BinaryCrossEntropyLoss` | `nn.BCELoss` |
| `CosineEmbeddingLoss` | `nn.CosineEmbeddingLoss` |
| `CrossEntropyLoss` | `nn.CrossEntropyLoss` |
| `CTCLoss` | `nn.CTCLoss` |
| `GramMatrixLoss` | _No direct equivalent_ |
| `HuberLoss` | `nn.HuberLoss` |
| `KLDivLoss` | `nn.KLDivLoss` |
| `LpLoss` | _No direct equivalent_ |
| `MseLoss` | `nn.MSELoss` |
| `PoissonNllLoss` | `nn.PoissonNLLLoss` |
| `RNNTLoss` | `torchaudio.functional.rnnt_loss` |
| `SmoothL1Loss` | `nn.SmoothL1Loss` |
burn-book/src/building-blocks/record.md
+85
View File
@@ -0,0 +1,85 @@
# Record
Records are how training state is saved and loaded with Burn. A record holds plain tensor data
(decoupled from the backend in use), so weights saved with one backend can be loaded on another, and
parameter initialization stays lazy.
All records serialize to the **burnpack** format (`.bpk`), Burn's compact binary container
implemented by the `burn-pack` crate.
## The burnpack format
A burnpack file has three parts:
- a small fixed-size **header** (a `"BURN"` magic number, a format version, and the metadata length);
- a **metadata** blob (CBOR) describing each tensor (name, dtype, shape, data offsets, optional
parameter id), any named **typed scalars**, and user key/value pairs;
- a **tensor data section** where each tensor's bytes start on a 256-byte boundary, so the data can
be read back with zero-copy / memory-mapped loading.
Storing typed scalars (integers, floats, booleans) alongside tensors is what lets the optimizer and
learning rate scheduler persist their non-tensor state in the same format.
## The three record types
| Record | Holds | Produced from |
| -------------------- | ------------------------------ | ---------------------------------------------- |
| `ModuleRecord` | a module's parameters | `module.into_record()` |
| `OptimizerRecord` | the optimizer state | `optimizer.to_record()` |
| `LrSchedulerRecord` | the learning rate scheduler | `scheduler.to_record()` |
Each record can be written to a file (`save` / `load`, which appends the `.bpk` extension when the
path has none) or to an in-memory byte buffer (`into_bytes` / `from_bytes`, useful for `no-std`
deployment where the bytes are embedded with the compiled code).
### `ModuleRecord`
`ModuleRecord` (in `burn::store`) holds a module's parameters keyed by their path within the module.
It is produced and applied through the `Module` trait itself:
```rust, ignore
use burn::store::ModuleRecord;
// Take a record and save it.
model.into_record().save("model")?; // writes model.bpk
// Load it back and apply it to an initialized module.
let record = ModuleRecord::load("model")?;
let model = ModelConfig::new().init(&device).load_record(record);
```
Load-time behavior is configured with builder methods on the record (ignored when saving):
- `.allow_partial(true)` — load even when some module parameters are absent from the record;
- `.validate(false)` — skip shape-mismatch / missing-tensor validation;
- `.cast_to_module_dtype()` / `.with_dtype_policy(..)` — cast the record's data to the module
parameter dtypes on load (by default the parameter adopts the record's dtype).
The save-side dtype is not configurable: the record stores whatever dtype the module currently holds.
To control the dtype applied on load, use `.cast_to_module_dtype()` / `.with_dtype_policy(..)` above.
Use `try_load_record` for the fallible variant of `load_record`.
### `OptimizerRecord` and `LrSchedulerRecord`
The optimizer and learning rate scheduler expose the same shape of API, used to checkpoint and resume
training:
```rust, ignore
// Optimizer state (no device needed on load; state migrates to each parameter's device on the
// next step).
optimizer.save("optim")?;
let optimizer = optimizer.load("optim")?;
// Learning rate scheduler state (scalars only).
scheduler.to_record().save("scheduler")?;
let scheduler = scheduler.load_record(LrSchedulerRecord::load("scheduler")?);
```
When training with the `Learner`, these records are saved and restored for you by the checkpointer —
see [Learner](./learner.md).
## Cross-framework formats
To import weights from other ecosystems (PyTorch `.pt`, SafeTensors) or to use the more advanced
store features (key remapping, filtering, half-precision storage), use the `burn-store` crate. See
[Saving and Loading Models](../saving-and-loading.md) for examples.
burn-book/src/building-blocks/tensor.md
+588
View File
@@ -0,0 +1,588 @@
# Tensor
As previously explained in the [model section](../basic-workflow/model.md), the Tensor struct has 3
generic arguments: the backend B, the dimensionality D, and the data type.
```rust, ignore
Tensor<B, D> // Float tensor (default)
Tensor<B, D, Float> // Explicit float tensor
Tensor<B, D, Int> // Int tensor
Tensor<B, D, Bool> // Bool tensor
```
Note that the specific element types used for `Float`, `Int`, and `Bool` tensors are defined by
backend implementations.
Burn Tensors are defined by the number of dimensions D in its declaration as opposed to its shape.
The actual shape of the tensor is inferred from its initialization. For example, a Tensor of size
(5,) is initialized as below:
```rust, ignore
let floats = [1.0, 2.0, 3.0, 4.0, 5.0];
// Get the default device
let device = Default::default();
// correct: Tensor is 1-Dimensional with 5 elements
let tensor_1 = Tensor::<Backend, 1>::from_floats(floats, &device);
// incorrect: let tensor_1 = Tensor::<Backend, 5>::from_floats(floats, &device);
// this will lead to an error and is for creating a 5-D tensor
```
### Initialization
Burn Tensors are primarily initialized using the `from_data()` method which takes the `TensorData`
struct as input. The `TensorData` struct has two public fields: `shape` and `dtype`. The `value`,
now stored as bytes, is private but can be accessed via any of the following methods: `as_slice`,
`as_mut_slice`, `to_vec` and `iter`. To retrieve the data from a tensor, the method `.to_data()`
should be employed when intending to reuse the tensor afterward. Alternatively, `.into_data()` is
recommended for one-time use. Let's look at a couple of examples for initializing a tensor from
different inputs.
```rust, ignore
// Initialization from a given Backend (Wgpu)
let tensor_1 = Tensor::<Wgpu, 1>::from_data([1.0, 2.0, 3.0], &device);
// Initialization from a generic Backend
let tensor_2 = Tensor::<Backend, 1>::from_data(TensorData::from([1.0, 2.0, 3.0]), &device);
// Initialization using from_floats (Recommended for f32 ElementType)
// Will be converted to TensorData internally.
let tensor_3 = Tensor::<Backend, 1>::from_floats([1.0, 2.0, 3.0], &device);
// Initialization of Int Tensor from array slices
let arr: [i32; 6] = [1, 2, 3, 4, 5, 6];
let tensor_4 = Tensor::<Backend, 1, Int>::from_data(TensorData::from(&arr[0..3]), &device);
// Initialization from a custom type
struct BodyMetrics {
age: i8,
height: i16,
weight: f32
}
let bmi = BodyMetrics{
age: 25,
height: 180,
weight: 80.0
};
let data = TensorData::from([bmi.age as f32, bmi.height as f32, bmi.weight]);
let tensor_5 = Tensor::<Backend, 1>::from_data(data, &device);
```
## Ownership and Cloning
Almost all Burn operations take ownership of the input tensors. Therefore, reusing a tensor multiple
times will necessitate cloning it. Let's look at an example to understand the ownership rules and
cloning better. Suppose we want to do a simple min-max normalization of an input tensor.
```rust, ignore
let input = Tensor::<Wgpu, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &device);
let min = input.min();
let max = input.max();
let input = (input - min).div(max - min);
```
With PyTorch tensors, the above code would work as expected. However, Rust's strict ownership rules
will give an error and prevent using the input tensor after the first `.min()` operation. The
ownership of the input tensor is transferred to the variable `min` and the input tensor is no longer
available for further operations. Burn Tensors like most complex primitives do not implement the
`Copy` trait and therefore have to be cloned explicitly. Now let's rewrite a working example of
doing min-max normalization with cloning.
```rust, ignore
let input = Tensor::<Wgpu, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &device);
let min = input.clone().min();
let max = input.clone().max();
let input = (input.clone() - min.clone()).div(max - min);
println!("{}", input.to_data());// Success: [0.0, 0.33333334, 0.6666667, 1.0]
// Notice that max, min have been moved in last operation so
// the below print will give an error.
// If we want to use them for further operations,
// they will need to be cloned in similar fashion.
// println!("{:?}", min.to_data());
```
We don't need to be worried about memory overhead because with cloning, the tensor's buffer isn't
copied, and only a reference to it is increased. This makes it possible to determine exactly how
many times a tensor is used, which is very convenient for reusing tensor buffers or even fusing
operations into a single kernel ([burn-fusion](https://burn.dev/docs/burn_fusion/index.htmls)). For
that reason, we don't provide explicit inplace operations. If a tensor is used only one time,
inplace operations will always be used when available.
## Tensor Operations
Normally with PyTorch, explicit inplace operations aren't supported during the backward pass, making
them useful only for data preprocessing or inference-only model implementations. With Burn, you can
focus more on _what_ the model should do, rather than on _how_ to do it. We take the responsibility
of making your code run as fast as possible during training as well as inference. The same
principles apply to broadcasting; all operations support broadcasting unless specified otherwise.
Here, we provide a list of all supported operations along with their PyTorch equivalents. Note that
for the sake of simplicity, we ignore type signatures. For more details, refer to the
[full documentation](https://docs.rs/burn/latest/burn/tensor/struct.Tensor.html).
### Basic Operations
Those operations are available for all tensor kinds: `Int`, `Float`, and `Bool`.
| Burn | PyTorch Equivalent |
| ---------------------------------------------------- | ------------------------------------------------------------------------- |
| `Tensor::cat(tensors, dim)` | `torch.cat(tensors, dim)` |
| `Tensor::empty(shape, options)` | `torch.empty(shape, device=device, dtype=dtype)` |
| `Tensor::from_primitive(primitive)` | N/A |
| `Tensor::stack(tensors, dim)` | `torch.stack(tensors, dim)` |
| `tensor.all()` | `tensor.all()` |
| `tensor.all_dim(dim)` | `tensor.all(dim)` |
| `tensor.any()` | `tensor.any()` |
| `tensor.any_dim(dim)` | `tensor.any(dim)` |
| `tensor.chunk(num_chunks, dim)` | `tensor.chunk(num_chunks, dim)` |
| `tensor.split(split_size, dim)` | `tensor.split(split_size, dim)` |
| `tensor.split_with_sizes(split_sizes, dim)` | `tensor.split([split_sizes], dim)` |
| `tensor.device()` | `tensor.device` |
| `tensor.dtype()` | `tensor.dtype` |
| `tensor.dims()` | `tensor.size()` |
| `tensor.equal(other)` | `x == y` |
| `tensor.equal_elem(other)` | `tensor.eq(other)` |
| `tensor.expand(shape)` | `tensor.expand(shape)` |
| `tensor.flatten(start_dim, end_dim)` | `tensor.flatten(start_dim, end_dim)` |
| `tensor.flip(axes)` | `tensor.flip(axes)` |
| `tensor.full_like(fill_value)` | `torch.full_like(tensor, fill_value)` |
| `tensor.gather(dim, indices)` | `torch.gather(tensor, dim, indices)` |
| `tensor.into_data()` | N/A |
| `tensor.into_primitive()` | N/A |
| `tensor.into_scalar()` | `tensor.item()` |
| `tensor.mask_fill(mask, value)` | `tensor.masked_fill(mask, value)` |
| `tensor.mask_where(mask, value_tensor)` | `torch.where(mask, value_tensor, tensor)` |
| `tensor.movedim(src, dst)` | `tensor.movedim(src, dst)` |
| `tensor.narrow(dim, start, length)` | `tensor.narrow(dim, start, length)` |
| `tensor.not_equal(other)` | `x != y` |
| `tensor.not_equal_elem(scalar)` | `tensor.ne(scalar)` |
| `tensor.ones_like()` | `torch.ones_like(tensor)` |
| `tensor.permute(axes)` | `tensor.permute(axes)` |
| `tensor.repeat_dim(dim, times)` | `tensor.repeat(*[times if i == dim else 1 for i in range(tensor.dim())])` |
| `tensor.repeat(sizes)` | `tensor.repeat(sizes)` |
| `tensor.reshape(shape)` | `tensor.view(shape)` |
| `tensor.roll(shifts, dims)` | `tensor.roll(shifts, dims)` |
| `tensor.roll_dim(shift, dim)` | `tensor.roll([shift], [dim])` |
| `tensor.scatter(dim, indices, values, update)` | `tensor.scatter_add(dim, indices, values)` |
| `tensor.scatter_nd(indices, values, update)` | N/A |
| `tensor.gather_nd(indices)` | N/A |
| `tensor.select(dim, indices)` | `tensor.index_select(dim, indices)` |
| `tensor.select_assign(dim, indices, values, update)` | `tensor.index_add(dim, indices, values)` |
| `tensor.shape()` | `tensor.shape` |
| `tensor.slice(slices)` | `tensor[(*ranges,)]` |
| `tensor.slice_assign(slices, values)` | `tensor[(*ranges,)] = values` |
| `tensor.slice_fill(slices, value)` | `tensor[(*ranges,)] = value` |
| `tensor.slice_dim(dim, slice)` | N/A |
| `tensor.squeeze()` | `tensor.squeeze()` |
| `tensor.squeeze_dim(dim)` | `tensor.squeeze(dim)` |
| `tensor.squeeze_dims(dims)` | `tensor.squeeze(dims)` where `dims` is a tuple of ints |
| `tensor.swap_dims(dim1, dim2)` | `tensor.transpose(dim1, dim2)` |
| `tensor.take(dim, indices)` | `numpy.take(tensor, indices, dim)` |
| `tensor.to_data()` | N/A |
| `tensor.to_device(device)` | `tensor.to(device)` |
| `tensor.transpose()` | `tensor.T` |
| `tensor.t()` | `tensor.T` |
| `tensor.unsqueeze()` | N/A |
| `tensor.unsqueeze_dim(dim)` | `tensor.unsqueeze(dim)` |
| `tensor.unsqueeze_dims(dims)` | N/A |
| `tensor.zeros_like()` | `torch.zeros_like(tensor)` |
| `Tensor::full(shape, fill_value, options)` | `torch.full(shape, fill_value, device=device, dtype=dtype)` |
| `Tensor::ones(shape, options)` | `torch.ones(shape, device=device, dtype=dtype)` |
| `Tensor::zeros(shape, options)` | `torch.zeros(shape, device=device, dtype=dtype)` |
### Numeric Operations
Those operations are available for numeric tensor kinds: `Float` and `Int`.
| Burn | PyTorch Equivalent |
| --------------------------------------------------------------- | --------------------------------------------- |
| `tensor.abs()` | `torch.abs(tensor)` |
| `tensor.add(other)` or `tensor + other` | `tensor + other` |
| `tensor.add_scalar(scalar)` or `tensor + scalar` | `tensor + scalar` |
| `tensor.all_close(other, atol, rtol)` | `torch.allclose(tensor, other, atol, rtol)` |
| `tensor.argmax(dim)` | `tensor.argmax(dim)` |
| `tensor.argmin(dim)` | `tensor.argmin(dim)` |
| `tensor.argsort(dim)` | `tensor.argsort(dim)` |
| `tensor.argsort_descending(dim)` | `tensor.argsort(dim, descending=True)` |
| `tensor.bool()` | `tensor.bool()` |
| `tensor.clamp(min, max)` | `torch.clamp(tensor, min=min, max=max)` |
| `tensor.clamp_max(max)` | `torch.clamp(tensor, max=max)` |
| `tensor.clamp_min(min)` | `torch.clamp(tensor, min=min)` |
| `tensor.cumsum(dim)` | `tensor.cumsum(dim)` |
| `tensor.cumprod(dim)` | `tensor.cumprod(dim)` |
| `tensor.cummin(dim)` | `tensor.cummin(dim)` |
| `tensor.cummax(dim)` | `tensor.cummax(dim)` |
| `tensor.div(other)` or `tensor / other` | `tensor / other` |
| `tensor.div_scalar(scalar)` or `tensor / scalar` | `tensor / scalar` |
| `tensor.dot(other)` | `torch.dot(tensor, other)` |
| `tensor.greater(other)` | `tensor.gt(other)` |
| `tensor.greater_elem(scalar)` | `tensor.gt(scalar)` |
| `tensor.greater_equal(other)` | `tensor.ge(other)` |
| `tensor.greater_equal_elem(scalar)` | `tensor.ge(scalar)` |
| `tensor.lower(other)` | `tensor.lt(other)` |
| `tensor.lower_elem(scalar)` | `tensor.lt(scalar)` |
| `tensor.lower_equal(other)` | `tensor.le(other)` |
| `tensor.lower_equal_elem(scalar)` | `tensor.le(scalar)` |
| `tensor.max()` | `tensor.max()` |
| `tensor.max_abs()` | `tensor.abs().max()` |
| `tensor.max_abs_dim(dim)` | `tensor.abs().max(dim, keepdim=True)` |
| `tensor.max_abs_dims(dims)` | `tensor.abs().max(dims, keepdim=True)` |
| `tensor.max_dim(dim)` | `tensor.max(dim, keepdim=True)` |
| `tensor.max_dims(dims)` | `tensor.max(dims, keepdim=True)` |
| `tensor.max_dim_with_indices(dim)` | N/A |
| `tensor.max_pair(other)` | `torch.Tensor.max(a,b)` |
| `tensor.mean()` | `tensor.mean()` |
| `tensor.mean_dim(dim)` | `tensor.mean(dim, keepdim=True)` |
| `tensor.mean_dims(dims)` | `tensor.mean(dims, keepdim=True)` |
| `tensor.min()` | `tensor.min()` |
| `tensor.min_dim(dim)` | `tensor.min(dim, keepdim=True)` |
| `tensor.min_dims(dims)` | `tensor.min(dims, keepdim=True)` |
| `tensor.min_dim_with_indices(dim)` | N/A |
| `tensor.min_pair(other)` | `torch.Tensor.min(a,b)` |
| `tensor.mul(other)` or `tensor * other` | `tensor * other` |
| `tensor.mul_scalar(scalar)` or `tensor * scalar` | `tensor * scalar` |
| `tensor.neg()` or `-tensor` | `-tensor` |
| `tensor.one_hot(num_classes)` | `torch.nn.functional.one_hot` |
| `tensor.one_hot_fill(num_classes, on_value, off_value, axis)` | N/A |
| `tensor.pad(pads, mode)` | `torch.nn.functional.pad(tensor, pads, mode)` |
| `tensor.powf(other)` or `tensor.powi(intother)` | `tensor.pow(other)` |
| `tensor.powf_scalar(scalar)` or `tensor.powi_scalar(intscalar)` | `tensor.pow(scalar)` |
| `tensor.prod()` | `tensor.prod()` |
| `tensor.prod_dim(dim)` | `tensor.prod(dim, keepdim=True)` |
| `tensor.prod_dims(dims)` | `tensor.prod(dims, keepdim=True)` |
| `tensor.rem(other)` or `tensor % other` | `tensor % other` |
| `tensor.sign()` | `tensor.sign()` |
| `tensor.sort(dim)` | `tensor.sort(dim).values` |
| `tensor.sort_descending(dim)` | `tensor.sort(dim, descending=True).values` |
| `tensor.sort_descending_with_indices(dim)` | `tensor.sort(dim, descending=True)` |
| `tensor.sort_with_indices(dim)` | `tensor.sort(dim)` |
| `tensor.sub(other)` or `tensor - other` | `tensor - other` |
| `tensor.sub_scalar(scalar)` or `tensor - scalar` | `tensor - scalar` |
| `tensor.sum()` | `tensor.sum()` |
| `tensor.sum_dim(dim)` | `tensor.sum(dim, keepdim=True)` |
| `tensor.sum_dims(dims)` | `tensor.sum(dims, keepdim=True)` |
| `tensor.sum_dims_squeeze(dims)` | `tensor.sum(dims, keepdim=False)` |
| `tensor.topk(k, dim)` | `tensor.topk(k, dim).values` |
| `tensor.topk_with_indices(k, dim)` | `tensor.topk(k, dim)` |
| `tensor.tril(diagonal)` | `torch.tril(tensor, diagonal)` |
| `tensor.triu(diagonal)` | `torch.triu(tensor, diagonal)` |
| `tensor.unfold(dim, size, step)` | `tensor.unfold(dim, size, step)` |
| `Tensor::eye(size, device)` | `torch.eye(size, device=device)` |
| `scalar - tensor` | `scalar - tensor` |
### Float Operations
Those operations are only available for `Float` tensors.
| Burn API | PyTorch Equivalent |
| -------------------------------------------- | ------------------------------------------ |
| `tensor.acos()` | `tensor.acos()` |
| `tensor.acosh()` | `tensor.acosh()` |
| `tensor.asin()` | `tensor.asin()` |
| `tensor.asinh()` | `tensor.asinh()` |
| `tensor.atan()` | `tensor.atan()` |
| `tensor.atanh()` | `tensor.atanh()` |
| `tensor.atan2(other_tensor)` | `tensor.atan2(other_tensor)` |
| `tensor.cast(dtype)` | `tensor.to(dtype)` |
| `tensor.ceil()` | `tensor.ceil()` |
| `tensor.contains_nan()` | N/A |
| `tensor.cos()` | `tensor.cos()` |
| `tensor.cosh()` | `tensor.cosh()` |
| `tensor.cross(other)` | `torch.cross(tensor, other)` |
| `tensor.deg2rad()` | `torch.deg2rad()` |
| `tensor.erf()` | `tensor.erf()` |
| `tensor.exp()` | `tensor.exp()` |
| `tensor.floor()` | `tensor.floor()` |
| `tensor.fmod(other)` | `tensor.fmod(other)` |
| `tensor.fmod_scalar(scalar)` | `tensor.fmod(scalar)` |
| `tensor.from_floats(floats, device)` | N/A |
| `tensor.int()` | Similar to `tensor.to(torch.long)` |
| `tensor.is_close(other, atol, rtol)` | `torch.isclose(tensor, other, atol, rtol)` |
| `tensor.is_finite()` | `torch.isfinite(tensor)` |
| `tensor.is_inf()` | `torch.isinf(tensor)` |
| `tensor.is_nan()` | `torch.isnan(tensor)` |
| `tensor.log()` | `tensor.log()` |
| `tensor.log1p()` | `tensor.log1p()` |
| `tensor.matmul(other)` | `tensor.matmul(other)` |
| `tensor.rad2deg()` | `torch.rad2deg()` |
| `tensor.random(shape, distribution, device)` | N/A |
| `tensor.random_like(distribution)` | `torch.rand_like()` only uniform |
| `tensor.recip()` or `1.0 / tensor` | `tensor.reciprocal()` or `1.0 / tensor` |
| `tensor.round()` | `tensor.round()` |
| `tensor.sin()` | `tensor.sin()` |
| `tensor.sinh()` | `tensor.sinh()` |
| `tensor.square()` | `tensor.square()` |
| `tensor.sqrt()` | `tensor.sqrt()` |
| `tensor.tan()` | `tensor.tan()` |
| `tensor.tanh()` | `tensor.tanh()` |
| `tensor.trunc()` | `tensor.trunc()` |
| `tensor.var(dim)` | `tensor.var(dim)` |
| `tensor.var_bias(dim)` | N/A |
| `tensor.var_mean(dim)` | N/A |
| `tensor.var_mean_bias(dim)` | N/A |
| `tensor.median(dim)` | `tensor.median(dim)` |
| `tensor.median_with_indices(dim)` | `tensor.median(dim)` |
### Int Operations
Those operations are only available for `Int` tensors.
| Burn API | PyTorch Equivalent |
| ------------------------------------------------ | ------------------------------------------------------- |
| `Tensor::arange(5..10, device)` | `tensor.arange(start=5, end=10, device=device)` |
| `Tensor::arange_step(5..10, 2, device)` | `tensor.arange(start=5, end=10, step=2, device=device)` |
| `tensor.bitwise_and(other)` | `torch.bitwise_and(tensor, other)` |
| `tensor.bitwise_and_scalar(scalar)` | `torch.bitwise_and(tensor, scalar)` |
| `tensor.bitwise_not()` | `torch.bitwise_not(tensor)` |
| `tensor.bitwise_left_shift(other)` | `torch.bitwise_left_shift(tensor, other)` |
| `tensor.bitwise_left_shift_scalar(scalar)` | `torch.bitwise_left_shift(tensor, scalar)` |
| `tensor.bitwise_right_shift(other)` | `torch.bitwise_right_shift(tensor, other)` |
| `tensor.bitwise_right_shift_scalar(scalar)` | `torch.bitwise_right_shift(tensor, scalar)` |
| `tensor.bitwise_or(other)` | `torch.bitwise_or(tensor, other)` |
| `tensor.bitwise_or_scalar(scalar)` | `torch.bitwise_or(tensor, scalar)` |
| `tensor.bitwise_xor(other)` | `torch.bitwise_xor(tensor, other)` |
| `tensor.bitwise_xor_scalar(scalar)` | `torch.bitwise_xor(tensor, scalar)` |
| `tensor.float()` | `tensor.to(torch.float)` |
| `tensor.from_ints(ints)` | N/A |
| `tensor.cartesian_grid(shape, device)` | N/A |
### Bool Operations
Those operations are only available for `Bool` tensors.
| Burn API | PyTorch Equivalent |
| ------------------------------------ | ------------------------------- |
| `Tensor::diag_mask(shape, diagonal)` | N/A |
| `Tensor::tril_mask(shape, diagonal)` | N/A |
| `Tensor::triu_mask(shape, diagonal)` | N/A |
| `tensor.argwhere()` | `tensor.argwhere()` |
| `tensor.bool_and()` | `tensor.logical_and()` |
| `tensor.bool_not()` | `tensor.logical_not()` |
| `tensor.bool_or()` | `tensor.logical_or()` |
| `tensor.bool_xor()` | `tensor.logical_xor()` |
| `tensor.float()` | `tensor.to(torch.float)` |
| `tensor.int()` | `tensor.to(torch.long)` |
| `tensor.nonzero()` | `tensor.nonzero(as_tuple=True)` |
### Quantization Operations
Those operations are only available for `Float` tensors on backends that implement quantization
strategies.
| Burn API | PyTorch Equivalent |
| ---------------------------------- | ------------------ |
| `tensor.quantize(scheme, qparams)` | N/A |
| `tensor.dequantize()` | N/A |
## Activation Functions
| Burn API | PyTorch Equivalent |
| ------------------------------------------------ | -------------------------------------------------- |
| `activation::celu(tensor, alpha)` | `nn.functional.celu(tensor, alpha)` |
| `activation::elu(tensor, alpha)` | `nn.functional.elu(tensor, alpha)` |
| `activation::gelu(tensor)` | `nn.functional.gelu(tensor)` |
| `activation::glu(tensor, dim)` | `nn.functional.glu(tensor, dim)` |
| `activation::hard_shrink(tensor, lambda)` | `nn.functional.hardshrink(tensor, lambd)` |
| `activation::hard_sigmoid(tensor, alpha, beta)` | `nn.functional.hardsigmoid(tensor)` |
| `activation::hard_swish(tensor)` | `nn.functional.hardswish(tensor)` |
| `activation::leaky_relu(tensor, negative_slope)` | `nn.functional.leaky_relu(tensor, negative_slope)` |
| `activation::log_sigmoid(tensor)` | `nn.functional.log_sigmoid(tensor)` |
| `activation::log_softmax(tensor, dim)` | `nn.functional.log_softmax(tensor, dim)` |
| `activation::mish(tensor)` | `nn.functional.mish(tensor)` |
| `activation::prelu(tensor,alpha)` | `nn.functional.prelu(tensor,weight)` |
| `activation::quiet_softmax(tensor, dim)` | `nn.functional.quiet_softmax(tensor, dim)` |
| `activation::relu(tensor)` | `nn.functional.relu(tensor)` |
| `activation::shrink(tensor, lambda, bias)` | _No direct equivalent_ |
| `activation::soft_shrink(tensor, lambda)` | `nn.functional.softshrink(tensor, lambd)` |
| `activation::sigmoid(tensor)` | `nn.functional.sigmoid(tensor)` |
| `activation::selu(tensor)` | `nn.functional.selu(tensor)` |
| `activation::silu(tensor)` | `nn.functional.silu(tensor)` |
| `activation::softmax(tensor, dim)` | `nn.functional.softmax(tensor, dim)` |
| `activation::softmin(tensor, dim)` | `nn.functional.softmin(tensor, dim)` |
| `activation::softplus(tensor, beta)` | `nn.functional.softplus(tensor, beta)` |
| `activation::softsign(tensor)` | `nn.functional.softsign(tensor)` |
| `activation::tanh(tensor)` | `nn.functional.tanh(tensor)` |
| `activation::thresholded_relu(tensor, alpha)` | `nn.functional.threshold(tensor, alpha, 0)` |
## Grid Functions
| Burn API | PyTorch Equivalent |
| --------------------------------------------------- | -------------------------------------------------------------------- |
| `grid::affine_grid_2d(transformation_tensor, dims)` | `nn.functional.affine_grid(theta_tensor, size, align_corners)` |
| `grid::meshgrid(tensors, GridIndexing::Matrix)` | `torch.meshgrid(tensors, indexing="ij")` |
| `grid::meshgrid(tensors, GridIndexing::Cartesian)` | `torch.meshgrid(tensors, indexing="xy")` |
| `grid::meshgrid_stack(tensors, index_pos)` | _No direct equivalent_ |
## Linalg Functions
| Burn API | PyTorch Equivalent |
| -------------------------------------------------- | --------------------------------------------------- |
| `linalg::cosine_similarity(x1, x2, dim, eps)` | `nn.functional.cosine_similarity(x1, x2, dim, eps)` |
| `linalg::det(tensor)` | `torch.linalg.det(tensor)` |
| `linalg::diag(tensor)` | `torch.diag(tensor)` |
| `linalg::l0_norm(tensor, dim)` | _No direct equivalent_ |
| `linalg::l1_norm(tensor, dim)` | _No direct equivalent_ |
| `linalg::l2_norm(tensor, dim)` | _No direct equivalent_ |
| `linalg::lp_norm(tensor, p, dim)` | _No direct equivalent_ |
| `linalg::lu(tensor)` | `torch.linalg.lu(tensor)` |
| `linalg::matvec(matrix, vector)` | `torch.matmul(matrix, vector)` / `@` operator |
| `linalg::max_abs_norm(tensor, dim)` | _No direct equivalent_ |
| `linalg::min_abs_norm(tensor, dim)` | _No direct equivalent_ |
| `linalg::outer(lhs, rhs)` | `torch.outer(lhs, rhs)` / `einsum("bi,bj->bij", …)` |
| `linalg::outer_dim(lhs, rhs, dim)` | _No direct equivalent_ |
| `linalg::trace(tensor)` | `torch.trace(tensor)` |
| `linalg::vector_norm(tensor, p, dim)` | `torch.linalg.vector_norm(tensor, p, dim)` |
| `linalg::vector_normalize(tensor, norm, dim, eps)` | `nn.functional.normalize(tensor, p, dim, eps)` |
## Signal Processing Functions
Signal-processing helpers live in `burn::tensor::signal` and operate on real-valued float
tensors. FFT length `n` (and `n_fft` in STFT) must currently be a power of two: when `n` is
`Some(size)`, the input is truncated or zero-padded to `size` and the output has
`size / 2 + 1` frequency bins. Non-power-of-two sizes panic at the public API boundary;
general arbitrary-size DFT support (Bluestein's algorithm) is a tracked follow-up.
| Burn API | PyTorch Equivalent |
| ----------------------------------------------------- | --------------------------------------------------------------------------------- |
| `signal::rfft(tensor, dim, n)` | `torch.fft.rfft(tensor, n, dim)` |
| `signal::irfft(re, im, dim, n)` | `torch.fft.irfft(complex, n, dim)` |
| `signal::stft(signal, window, options)` | `torch.stft(signal, n_fft, hop_length, win_length, window, center)` |
| `signal::istft(stft_matrix, window, length, options)` | `torch.istft(stft_matrix, n_fft, hop_length, win_length, window, center, length)` |
| `signal::blackman_window(size, periodic, options)` | `torch.blackman_window(size, periodic)` |
| `signal::hamming_window(size, periodic, options)` | `torch.hamming_window(size, periodic)` |
| `signal::hann_window(size, periodic, options)` | `torch.hann_window(size, periodic)` |
`stft` and `istft` share a `StftOptions` struct with fields `n_fft`, `hop_length`,
`win_length`, `center`, and `onesided`. Use `StftOptions::new(n_fft)` for PyTorch-style
defaults (`hop_length = n_fft / 4`, `win_length = None`, `center = true`, `onesided = true`).
The option set is validated on entry to both `stft` and `istft`; `n_fft` must be a power of
two and `hop_length <= effective_win_length` (the COLA prerequisite for invertibility).
## Displaying Tensor Details
Burn provides flexible options for displaying tensor information, allowing you to control the level
of detail and formatting to suit your needs.
### Basic Display
To display a detailed view of a tensor, you can simply use Rust's `println!` or `format!` macros:
```rust, ignore
let tensor = Tensor::<Backend, 2>::full([2, 3], 0.123456789, &Default::default());
println!("{}", tensor);
```
This will output:
```
Tensor {
data:
[[0.12345679, 0.12345679, 0.12345679],
[0.12345679, 0.12345679, 0.12345679]],
shape: [2, 3],
device: Cpu,
backend: "flex",
kind: "Float",
dtype: "f32",
}
```
### Controlling Precision
You can control the number of decimal places displayed using Rust's formatting syntax:
```rust
println!("{:.2}", tensor);
```
Output:
```
Tensor {
data:
[[0.12, 0.12, 0.12],
[0.12, 0.12, 0.12]],
shape: [2, 3],
device: Cpu,
backend: "flex",
kind: "Float",
dtype: "f32",
}
```
### Global Print Options
For more fine-grained control over tensor printing, Burn provides a `PrintOptions` struct and a
`set_print_options` function:
```rust, ignore
use burn::tensor::{set_print_options, PrintOptions};
let print_options = PrintOptions {
precision: Some(2),
..Default::default()
};
set_print_options(print_options);
```
Options:
- `precision`: Number of decimal places for floating-point numbers (default: None)
- `threshold`: Maximum number of elements to display before summarizing (default: 1000)
- `edge_items`: Number of items to show at the beginning and end of each dimension when summarizing
(default: 3)
### Checking Tensor Closeness
Burn provides a utility function `check_closeness` to compare two tensors and assess their
similarity. This function is particularly useful for debugging and validating tensor operations,
especially when working with floating-point arithmetic where small numerical differences can
accumulate. It's also valuable when comparing model outputs during the process of importing models
from other frameworks, helping to ensure that the imported model produces results consistent with
the original.
Here's an example of how to use `check_closeness`:
```rust, ignore
use burn::tensor::{check_closeness, Tensor};
type B = burn::backend::Flex;
let device = Default::default();
let tensor1 = Tensor::<B, 1>::from_floats(
[1.0, 2.0, 3.0, 4.0, 5.0, 6.001, 7.002, 8.003, 9.004, 10.1],
&device,
);
let tensor2 = Tensor::<B, 1>::from_floats(
[1.0, 2.0, 3.0, 4.000, 5.0, 6.0, 7.001, 8.002, 9.003, 10.004],
&device,
);
check_closeness(&tensor1, &tensor2);
```
The `check_closeness` function compares the two input tensors element-wise, checking their
absolute differences against a range of epsilon values. It then prints a detailed report showing
the percentage of elements that are within each tolerance level.
The output provides a breakdown for different epsilon values, allowing you to assess the closeness
of the tensors at various precision levels. This is particularly helpful when dealing with
operations that may introduce small numerical discrepancies.
The function uses color-coded output to highlight the results:
- Green [PASS]: All elements are within the specified tolerance.
- Yellow [WARN]: Most elements (90% or more) are within tolerance.
- Red [FAIL]: Significant differences are detected.
This utility can be invaluable when implementing or debugging tensor operations, especially those
involving complex mathematical computations or when porting algorithms from other frameworks. It's
also an essential tool when verifying the accuracy of imported models, ensuring that the Burn
implementation produces results that closely match those of the original model.
burn-book/src/custom-training-loop.md
+278
View File
@@ -0,0 +1,278 @@
# Custom Training Loops
Even though Burn comes with a project dedicated to simplifying training, it doesn't mean that you
have to use it. Sometimes you may have special needs for your training, and it might be faster to
just reimplement the training loop yourself. Also, you may just prefer implementing your own
training loop instead of using a pre-built one in general.
Burn's got you covered!
We will start from the same example shown in the [basic workflow](./basic-workflow) section, but
without using the `Learner` struct.
```rust, ignore
#[derive(Config, Debug)]
pub struct MnistTrainingConfig {
#[config(default = 10)]
pub num_epochs: usize,
#[config(default = 64)]
pub batch_size: usize,
#[config(default = 4)]
pub num_workers: usize,
#[config(default = 42)]
pub seed: u64,
#[config(default = 1e-4)]
pub lr: f64,
pub model: ModelConfig,
pub optimizer: AdamConfig,
}
pub fn run<B: AutodiffBackend>(device: B::Device) {
// Create the configuration.
let config_model = ModelConfig::new(10, 1024);
let config_optimizer = AdamConfig::new();
let config = MnistTrainingConfig::new(config_model, config_optimizer);
B::seed(&device, config.seed);
// Create the model and optimizer.
let mut model = config.model.init::<B>(&device);
let mut optim = config.optimizer.init();
// Create the batcher.
let batcher = MnistBatcher::default();
// Create the dataloaders.
let dataloader_train = DataLoaderBuilder::new(batcher.clone())
.batch_size(config.batch_size)
.shuffle(config.seed)
.num_workers(config.num_workers)
.build(MnistDataset::train());
let dataloader_test = DataLoaderBuilder::new(batcher)
.batch_size(config.batch_size)
.shuffle(config.seed)
.num_workers(config.num_workers)
.build(MnistDataset::test());
...
}
```
As seen with the previous example, setting up the configurations and the dataloader hasn't changed.
Now, let's move forward and write our own training loop:
```rust, ignore
pub fn run<B: AutodiffBackend>(device: B::Device) {
...
// Iterate over our training and validation loop for X epochs.
for epoch in 1..config.num_epochs + 1 {
// Implement our training loop.
for (iteration, batch) in dataloader_train.iter().enumerate() {
let output = model.forward(batch.images);
let loss = CrossEntropyLoss::new(None, &output.device())
.forward(output.clone(), batch.targets.clone());
let accuracy = accuracy(output, batch.targets);
println!(
"[Train - Epoch {} - Iteration {}] Loss {:.3} | Accuracy {:.3} %",
epoch,
iteration,
loss.clone().into_scalar(),
accuracy,
);
// Gradients for the current backward pass
let grads = loss.backward();
// Gradients linked to each parameter of the model.
let grads = GradientsParams::from_grads(grads, &model);
// Update the model using the optimizer.
model = optim.step(config.lr, model, grads);
}
// Get the model without autodiff.
let model_valid = model.valid();
// Implement our validation loop.
for (iteration, batch) in dataloader_test.iter().enumerate() {
let output = model_valid.forward(batch.images);
let loss = CrossEntropyLoss::new(None, &output.device())
.forward(output.clone(), batch.targets.clone());
let accuracy = accuracy(output, batch.targets);
println!(
"[Valid - Epoch {} - Iteration {}] Loss {} | Accuracy {}",
epoch,
iteration,
loss.clone().into_scalar(),
accuracy,
);
}
}
}
```
In the previous code snippet, we can observe that the loop starts from epoch `1` and goes up to
`num_epochs`. Within each epoch, we iterate over the training dataloader. During this process, we
execute the forward pass, which is necessary for computing both the loss and accuracy. To maintain
simplicity, we print the results to stdout.
Upon obtaining the loss, we can invoke the `backward()` function, which returns the gradients
specific to each variable. It's important to note that we need to map these gradients to their
corresponding parameters using the `GradientsParams` type. This step is essential because you might
run multiple different autodiff graphs and accumulate gradients for each parameter id.
Finally, we can perform the optimization step using the learning rate, the model, and the computed
gradients. It's worth mentioning that, unlike PyTorch, there's no need to register the gradients
with the optimizer, nor do you have to call `zero_grad`. The gradients are automatically consumed
during the optimization step. If you're interested in gradient accumulation, you can easily achieve
this by using the `GradientsAccumulator`.
```rust, ignore
let mut accumulator = GradientsAccumulator::new();
let grads = model.backward();
let grads = GradientsParams::from_grads(grads, &model);
accumulator.accumulate(&model, grads); ...
let grads = accumulator.grads(); // Pop the accumulated gradients.
```
Note that after each epoch, we include a validation loop to assess our model's performance on
previously unseen data. To disable gradient tracking during this validation step, we can invoke
`model.valid()`, which provides a model on the inner backend without autodiff capabilities. It's
important to emphasize that we've declared our validation batcher to be on the inner backend,
specifically `MnistBatcher<B::InnerBackend>`; not using `model.valid()` will result in a compilation
error.
You can find the code above available as an
[example](https://github.com/tracel-ai/burn/tree/main/examples/custom-training-loop) for you to
test.
## Multiple optimizers
It's common practice to set different learning rates, optimizer parameters, or use different optimizers entirely, for different parts
of a model. In Burn, each `GradientParams` can contain only a subset of gradients to actually apply with an optimizer.
This allows you to flexibly mix and match optimizers!
```rust,ignore
// Start with calculating all gradients
let grads = loss.backward();
// Now split the gradients into various parts.
let grads_conv1 = GradientParams::from_module(&mut grads, &model.conv1);
let grads_conv2 = GradientParams::from_module(&mut grads, &model.conv2);
// You can step the model with these gradients, using different learning
// rates for each param. You could also use an entirely different optimizer here!
model = optim.step(config.lr * 2.0, model, grads_conv1);
model = optim.step(config.lr * 4.0, model, grads_conv2);
// For even more granular control you can split off individual parameter
// eg. a linear bias usually needs a smaller learning rate.
if let Some(bias) == model.linear1.bias {
let grads_bias = GradientParams::from_params(&mut grads, &model.linear1, &[bias.id]);
model = optim.step(config.lr * 0.1, model, grads_bias);
}
// Note that above calls remove gradients, so we can just get all "remaining" gradients.
let grads = GradientsParams::from_grads(grads, &model);
model = optim.step(config.lr, model, grads);
```
## Custom Type
The explanations above demonstrate how to create a basic training loop. However, you may find it
beneficial to organize your program using intermediary types. There are various ways to do this, but
it requires getting comfortable with generics.
If you wish to group the optimizer and the model into the same structure, you have several options.
It's important to note that the optimizer trait depends on both the `AutodiffModule` trait and the
`AutodiffBackend` trait, while the module only depends on the `AutodiffBackend` trait.
Here's a closer look at how you can create your types:
**Create a struct that is generic over the backend and the optimizer, with a predefined model.**
```rust, ignore
struct Learner<B, O>
where
B: AutodiffBackend,
{
model: Model<B>,
optim: O,
}
```
This is quite straightforward. You can be generic over the backend since it's used with the concrete
type `Model` in this case.
**Create a struct that is generic over the model and the optimizer.**
```rust, ignore
struct Learner<M, O> {
model: M,
optim: O,
}
```
This option is a quite intuitive way to declare the struct. You don't need to write type constraints
with a `where` statement when defining a struct; you can wait until you implement the actual
function. However, with this struct, you may encounter some issues when trying to implement code
blocks to your struct.
```rust, ignore
impl<B, M, O> Learner<M, O>
where
B: AutodiffBackend,
M: AutodiffModule<B>,
O: Optimizer<M, B>,
{
pub fn step(&mut self, _batch: MnistBatch<B>) {
//
}
}
```
This will result in the following compilation error:
```console
1. the type parameter `B` is not constrained by the impl trait, self type, or predicates
unconstrained type parameter [E0207]
```
To resolve this issue, you have two options. The first one is to make your function generic over
the backend and add your trait constraint within its definition:
```rust, ignore
#[allow(dead_code)]
impl<M, O> Learner2<M, O> {
pub fn step<B: AutodiffBackend>(&mut self, _batch: MnistBatch<B>)
where
B: AutodiffBackend,
M: AutodiffModule<B>,
O: Optimizer<M, B>,
{
//
}
}
```
However, some people may prefer to have the constraints on the implementation block itself. In that
case, you can make your struct generic over the backend using `PhantomData<B>`.
**Create a struct that is generic over the backend, the model, and the optimizer.**
```rust, ignore
struct Learner3<B, M, O> {
model: M,
optim: O,
_b: PhantomData<B>,
}
```
You might wonder why `PhantomData` is required. Each generic argument must be used as a field when
declaring a struct. When you don't need the generic argument, you can use `PhantomData` to mark it
as a zero sized type.
These are just some suggestions on how to define your own types, but you are free to use any pattern
that you prefer.
burn-book/src/distributed-computing.md
+1
View File
@@ -0,0 +1 @@
# Distributed Computing
burn-book/src/examples.md
+101
View File
@@ -0,0 +1,101 @@
# Examples
In the [next chapter](./basic-workflow) you'll have the opportunity to implement the whole Burn
`guide` example yourself in a step by step manner.
Many additional Burn examples are available in the
[examples](https://github.com/tracel-ai/burn/tree/main/examples) directory. Burn examples are
organized as library crates with one or more examples that are executable binaries. An example can
then be executed using the following cargo command line in the root of the Burn repository:
```bash
cargo run --example <example name>
```
To learn more about crates and examples, read the Rust section below.
<details>
<summary><strong>🦀 About Rust crates</strong></summary>
Each Burn example is a **package** which are subdirectories of the `examples` directory. A package
is composed of one or more **crates**.
A package is a bundle of one or more crates that provides a set of functionality. A package contains
a `Cargo.toml` file that describes how to build those crates.
A crate is a compilation unit in Rust. It could be a single file, but it is often easier to split up
crates into multiple **modules**.
A module lets us organize code within a crate for readability and easy reuse. Modules also allow us
to control the _privacy_ of items. For instance the `pub(crate)` keyword is employed to make a
module publicly available inside the crate. In the snippet below there are four modules declared,
two of them are public and visible to the users of the crates, one of them is public inside the
crate only and crate users cannot see it, at last one is private when there is no keyword. These
modules can be single files or a directory with a `mod.rs` file inside.
```rust, ignore
pub mod data;
pub mod inference;
pub(crate) mod model;
mod training;
```
A crate can come in one of two forms: a **binary crate** or a **library crate**. When compiling a
crate, the compiler first looks in the crate root file (`src/lib.rs` for a library crate and
`src/main.rs` for a binary crate). Any module declared in the crate root file will be inserted in
the crate for compilation.
All Burn examples are library crates and they can contain one or more executable examples that uses
the library. We even have some Burn examples that uses the library crate of other examples.
The examples are unique files under the `examples` directory. Each file produces an executable file
with the same name. Each example can then be executed with `cargo run --example <executable name>`.
Below is a file tree of a typical Burn example package:
```
examples/burn-example
├── Cargo.toml
├── examples
│ ├── example1.rs ---> compiled to example1 binary
│ ├── example2.rs ---> compiled to example2 binary
│ └── ...
└── src
├── lib.rs ---> this is the root file for a library
├── module1.rs
├── module2.rs
└── ...
```
</details><br>
The following additional examples are currently available if you want to check them out:
| Example | Description |
| :-------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| [Custom CSV Dataset](https://github.com/tracel-ai/burn/tree/main/examples/custom-csv-dataset) | Implements a dataset to parse CSV data for a regression task. |
| [Regression](https://github.com/tracel-ai/burn/tree/main/examples/simple-regression) | Trains a simple MLP on the California Housing dataset to predict the median house value for a district. |
| [Custom Image Dataset](https://github.com/tracel-ai/burn/tree/main/examples/custom-image-dataset) | Trains a simple CNN on custom image dataset following a simple folder structure. |
| [Custom Renderer](https://github.com/tracel-ai/burn/tree/main/examples/custom-renderer) | Implements a custom renderer to display the [`Learner`](./building-blocks/learner.md) progress. |
| [Image Classification Web](https://github.com/tracel-ai/burn-onnx/tree/main/examples/image-classification-web) | Image classification web browser demo using Burn, WGPU and WebAssembly. |
| [MNIST Inference on Web](https://github.com/tracel-ai/burn/tree/main/examples/mnist-inference-web) | An interactive MNIST inference demo in the browser. The demo is available [online](https://burn.dev/demo/). |
| [MNIST Training](https://github.com/tracel-ai/burn/tree/main/examples/mnist) | Demonstrates how to train a custom [`Module`](./building-blocks/module.md) (MLP) with the [`Learner`](./building-blocks/learner.md) configured to log metrics and keep training checkpoints. |
| [ONNX Import Inference](https://github.com/tracel-ai/burn-onnx/tree/main/examples/onnx-inference) | Imports an ONNX model pre-trained on MNIST to perform inference on a sample image with Burn. |
| [PyTorch Import Inference](https://github.com/tracel-ai/burn/tree/main/examples/import-model-weights) | Imports a PyTorch model pre-trained on MNIST to perform inference on a sample image with Burn. |
| [Text Classification](https://github.com/tracel-ai/burn/tree/main/examples/text-classification) | Trains a text classification transformer model on the AG News or DbPedia datasets. The trained model can then be used to classify a text sample. |
| [Text Generation](https://github.com/tracel-ai/burn/tree/main/examples/text-generation) | Trains a text generation transformer model on the DbPedia dataset. |
| [Wasserstein GAN MNIST](https://github.com/tracel-ai/burn/tree/main/examples/wgan) | Trains a WGAN model to generate new handwritten digits based on MNIST. |
For more information on each example, see their respective `README.md` file. Be sure to check out
the [examples](https://github.com/tracel-ai/burn/tree/main/examples) directory for an up-to-date
list.
<div class="warning">
Note that some examples use the
[`datasets` library by HuggingFace](https://huggingface.co/docs/datasets/index) to download the
datasets required in the examples. This is a Python library, which means that you will need to
install Python before running these examples. This requirement will be clearly indicated in the
example's README when applicable.
</div>
burn-book/src/getting-started.md
+219
View File
@@ -0,0 +1,219 @@
# Getting Started
Burn is a deep learning framework in the Rust programming language. Therefore, it goes without
saying that one must understand the basic notions of Rust. Reading the first chapters of the
[Rust Book](https://doc.rust-lang.org/book/) is recommended, but don't worry if you're just starting
out. We'll try to provide as much context and reference to external resources when required. Just
look out for the **🦀 Rust Note** indicators.
## Installing Rust
For installation instructions, please refer to the
[installation page](https://doc.rust-lang.org/book/ch01-01-installation.html). It explains in
details the most convenient way for you to install Rust on your computer, which is the very first
thing to do to start using Burn.
## Creating a Burn application
Once Rust is correctly installed, create a new Rust application by using Rust's build system and
package manager Cargo. It is automatically installed with Rust.
<details>
<summary><strong>🦀 Cargo Cheat Sheet</strong></summary>
[Cargo](https://doc.rust-lang.org/cargo/) is a very useful tool to manage Rust projects because it
handles a lot of tasks. More precisely, it is used to compile your code, download the
libraries/packages your code depends on, and build said libraries.
Below is a quick cheat sheet of the main `cargo` commands you might use throughout this guide.
| Command | Description |
| ------------------- | -------------------------------------------------------------------------------------------- |
| `cargo new` _path_ | Create a new Cargo package in the given directory. |
| `cargo add` _crate_ | Add dependencies to the Cargo.toml manifest file. |
| `cargo build` | Compile the local package and all of its dependencies (in debug mode, use `-r` for release). |
| `cargo check` | Check the local package for compilation errors (much faster). |
| `cargo run` | Run the local package binary. |
For more information, check out
[Hello, Cargo!](https://doc.rust-lang.org/book/ch01-03-hello-cargo.html) in the Rust Book.
</details><br>
In the directory of your choice, run the following:
```console
cargo new my_burn_app
```
This will initialize the `my_burn_app` project directory with a `Cargo.toml` file and a `src`
directory with an auto-generated `main.rs` file inside. Head inside the directory to check:
```console
cd my_burn_app
```
Then, add Burn as a dependency:
```console
cargo add burn --features wgpu
```
Finally, compile the local package by executing the following:
```console
cargo build
```
That's it, you're ready to start! You have a project configured with Burn and the WGPU backend,
which allows to execute low-level operations on any platform using the GPU.
<div class="warning">
When using one of the `wgpu` backends, you may encounter compilation errors related to recursive
type evaluation. This is due to complex type nesting within the `wgpu` dependency chain.
To resolve this issue, add the following line at the top of your `main.rs` or `lib.rs` file:
```rust
#![recursion_limit = "256"]
```
The default recursion limit (128) is often just below the required depth (typically 130-150) due to
deeply nested associated types and trait bounds.
</div>
## Writing a code snippet
The `src/main.rs` was automatically generated by Cargo, so let's replace its content with the
following:
```rust, ignore
use burn::tensor::Tensor;
use burn::backend::Wgpu;
// Type alias for the backend to use.
type Backend = Wgpu;
fn main() {
let device = Default::default();
// Creation of two tensors, the first with explicit values and the second one with ones, with the same shape as the first
let tensor_1 = Tensor::<Backend, 2>::from_data([[2., 3.], [4., 5.]], &device);
let tensor_2 = Tensor::<Backend, 2>::ones_like(&tensor_1);
// Print the element-wise addition (done with the WGPU backend) of the two tensors.
println!("{}", tensor_1 + tensor_2);
}
```
<details>
<summary><strong>🦀 Use Declarations</strong></summary>
To bring any of the Burn module or item into scope, a `use` declaration is added.
In the example above, we wanted bring the `Tensor` struct and `Wgpu` backend into scope with the
following:
```rust, ignore
use burn::tensor::Tensor;
use burn::backend::Wgpu;
```
This is pretty self-explanatory in this case. But, the same declaration could be written as a
shortcut to simultaneously binding of multiple paths with a common prefix:
```rust, ignore
use burn::{tensor::Tensor, backend::Wgpu};
```
In this example, the common prefix is pretty short and there are only two items to bind locally.
Therefore, the first usage with two `use` declarations might be preferred. But know that both
examples are valid. For more details on the `use` keyword, take a look at
[this section](https://doc.rust-lang.org/book/ch07-04-bringing-paths-into-scope-with-the-use-keyword.html)
of the Rust Book or the
[Rust reference](https://doc.rust-lang.org/reference/items/use-declarations.html).
</details><br>
<details>
<summary><strong>🦀 Generic Data Types</strong></summary>
If you're new to Rust, you're probably wondering why we had to use `Tensor::<Backend, 2>::...`.
That's because the `Tensor` struct is [generic](https://doc.rust-lang.org/book/ch10-01-syntax.html)
over multiple concrete data types. More specifically, a `Tensor` can be defined using three generic
parameters: the backend, the number of dimensions (rank) and the data type (defaults to `Float`).
Here, we only specify the backend and number of dimensions since a `Float` tensor is used by
default. For more details on the `Tensor` struct, take a look at
[this section](./building-blocks/tensor.md).
Most of the time when generics are involved, the compiler can infer the generic parameters
automatically. In this case, the compiler needs a little help. This can usually be done in one of
two ways: providing a type annotation or binding the generic parameter via the _turbofish_ `::<>`
syntax. In the example above we used the so-called _turbofish_ syntax, but we could have used type
annotations instead like this:
```rust, ignore
let tensor_1: Tensor<Backend, 2> = Tensor::from_data([[2., 3.], [4., 5.]]);
let tensor_2 = Tensor::ones_like(&tensor_1);
```
You probably noticed that we provided a type annotation for the first tensor only and yet this
example still works. That's because the compiler (correctly) inferred that `tensor_2` had the same
generic parameters. The same could have been done in the original example, but specifying the
parameters for both is more explicit.
</details><br>
By running `cargo run`, you should now see the result of the addition:
```console
Tensor {
data:
[[3.0, 4.0],
[5.0, 6.0]],
shape: [2, 2],
device: DefaultDevice,
backend: "wgpu",
kind: "Float",
dtype: "f32",
}
```
While the previous example is somewhat trivial, the upcoming basic workflow section will walk you
through a much more relevant example for deep learning applications.
## Using `prelude`
Burn comes with a variety of things in its core library. When creating a new model or using an
existing one for inference, you may need to import every single component you used, which could be a
little verbose.
To address it, a `prelude` module is provided, allowing you to easily import commonly used structs
and macros as a group:
```rust, ignore
use burn::prelude::*;
```
which is equal to:
```rust, ignore
use burn::{
config::Config,
module::Module,
nn,
tensor::{
backend::Backend, Bool, Device, ElementConversion, Float, Int, Shape, Tensor,
TensorData,
},
};
```
<div class="warning">
For the sake of simplicity, the subsequent chapters of this book will all use this form of importing
except in the [Building Blocks](./building-blocks) chapter, as explicit importing aids users in
grasping the usage of particular structures and macros.
</div>
burn-book/src/models-and-pretrained-weights.md
+30
View File
@@ -0,0 +1,30 @@
# Models and Pre-Trained Weights
## Models Repository
The [`models`](https://github.com/tracel-ai/models) repository contains definitions of different
deep learning models with examples for different domains like computer vision and natural language
processing.
This includes image classification models such as
[`MobileNetV2`](https://github.com/tracel-ai/models/tree/main/mobilenetv2-burn),
[`SqueezeNet`](https://github.com/tracel-ai/models/tree/main/squeezenet-burn) and
[`ResNet`](https://github.com/tracel-ai/models/tree/main/resnet-burn), object detection models such
as [`YOLOX`](https://github.com/tracel-ai/models/tree/main/yolox-burn) and language models like
[`BERT` and `RoBERTa`](https://github.com/tracel-ai/models/tree/main/bert-burn).
Be sure to check out the up-to-date
[collection of models](https://github.com/tracel-ai/models?tab=readme-ov-file#collection-of-official-models)
to get you started. Pre-trained weights are available for every supported architecture in this
collection. You will also find a spotlight of
[community contributed models](https://github.com/tracel-ai/models?tab=readme-ov-file#community-contributions).
## Burn-LM (alpha)
[`Burn-LM`](https://github.com/tracel-ai/burn-lm) is an LLM inference engine built on Burn. It
provides access to large language models with open-source pre-trained weights and supports running,
fine-tuning, and experimenting with them on any Burn backend.
Unlike tools focused solely on inference, Burn-LM is designed to work in a unified way across
different models and tasks, making it easier to explore both inference and training workflows within
the same framework.
burn-book/src/motivation.md
+28
View File
@@ -0,0 +1,28 @@
# Why Burn?
Why bother with the effort of creating an entirely new deep learning framework from scratch when
PyTorch, TensorFlow, and other frameworks already exist? Spoiler alert: Burn isn't merely a
replication of PyTorch or TensorFlow in Rust. It represents a novel approach, placing significant
emphasis on making the right compromises in the right areas to facilitate exceptional flexibility,
high performance, and a seamless developer experience. Burn isn’t a framework specialized for only
one type of application, it is designed to serve as a versatile framework suitable for a wide range
of research and production uses. The foundation of Burn's design revolves around three key user
profiles:
**Machine Learning Researchers** require tools to construct and execute experiments efficiently.
It’s essential for them to iterate quickly on their ideas and design testable experiments which can
help them discover new findings. The framework should facilitate the swift implementation of
cutting-edge research while ensuring fast execution for testing.
**Machine Learning Engineers** are another important demographic to keep in mind. Their focus leans
less on swift implementation and more on establishing robustness, seamless deployment, and
cost-effective operations. They seek dependable, economical models capable of achieving objectives
without excessive expense. The whole machine learning workflow —from training to inference— must be
as efficient as possible with minimal unpredictable behavior.
**Low level Software Engineers** working with hardware vendors want their processing units to run
models as fast as possible to gain competitive advantage. This endeavor involves harnessing
hardware-specific features such as Tensor Core for Nvidia. Since they are mostly working at a system
level, they want to have absolute control over how the computation will be executed.
The goal of Burn is to satisfy all of those personas!
burn-book/src/onnx-import.md
+278
View File
@@ -0,0 +1,278 @@
# ONNX Import
## Introduction
As deep learning evolves, interoperability between frameworks becomes crucial. Burn provides robust
support for importing [ONNX (Open Neural Network Exchange)](https://onnx.ai/onnx/intro/index.html)
models through the [`burn-onnx`](https://github.com/tracel-ai/burn-onnx) crate, enabling you to
leverage pre-trained models in your Rust-based deep learning projects.
## Why Import Models?
Importing pre-trained models offers several advantages:
1. **Time-saving**: Skip the resource-intensive process of training models from scratch.
2. **Access to state-of-the-art architectures**: Utilize cutting-edge models developed by
researchers and industry leaders.
3. **Transfer learning**: Fine-tune imported models for your specific tasks, benefiting from
knowledge transfer.
4. **Consistency across frameworks**: Maintain consistent performance when moving between
frameworks.
## Understanding ONNX
ONNX (Open Neural Network Exchange) is an open format designed to represent machine learning models
with these key features:
- **Framework agnostic**: Provides a common format that works across various deep learning
frameworks.
- **Comprehensive representation**: Captures both the model architecture and trained weights.
- **Wide support**: Compatible with popular frameworks like PyTorch, TensorFlow, and scikit-learn.
This standardization allows seamless movement of models between different frameworks and deployment
environments.
## Burn's ONNX Support
Burn's approach to ONNX import offers unique advantages:
1. **Native Rust code generation**: Translates ONNX models into Rust source code for deep
integration with Burn's ecosystem.
2. **Compile-time optimization**: Leverages the Rust compiler to optimize the generated code,
potentially improving performance.
3. **No runtime dependency**: Eliminates the need for an ONNX runtime, unlike many other solutions.
4. **Trainability**: Allows imported models to be further trained or fine-tuned using Burn.
5. **Portability**: Enables compilation for various targets, including WebAssembly and embedded
devices.
6. **Backend flexibility**: Works with any of Burn's supported backends.
## ONNX Compatibility
Burn recommends ONNX models use **opset version 16 or higher** for best compatibility. While models
with older opset versions may work, opset 16+ ensures access to all supported operators and their
latest behavior. If you encounter issues with an older model, consider upgrading it using the ONNX
version converter.
### Upgrading ONNX Models
There are two simple ways to upgrade your ONNX models to the recommended opset version:
Option 1: Use the provided utility script:
```
uv run --script https://raw.githubusercontent.com/tracel-ai/burn-onnx/refs/heads/main/onnx_opset_upgrade.py
```
Option 2: Use a custom Python script:
```python
import onnx
from onnx import version_converter, shape_inference
# Load your ONNX model
model = onnx.load('path/to/your/model.onnx')
# Convert the model to opset version 16
upgraded_model = version_converter.convert_version(model, 16)
# Apply shape inference to the upgraded model
inferred_model = shape_inference.infer_shapes(upgraded_model)
# Save the converted model
onnx.save(inferred_model, 'upgraded_model.onnx')
```
## Step-by-Step Guide
Follow these steps to import an ONNX model into your Burn project:
### Step 1: Update `Cargo.toml`
First, add the required dependencies to your `Cargo.toml`:
```toml
[dependencies]
burn = { version = "~0.21", features = ["flex"] }
[build-dependencies]
burn-onnx = "~0.21"
```
### Step 2: Update `build.rs`
In your `build.rs` file:
```rust, ignore
use burn_onnx::ModelGen;
fn main() {
ModelGen::new()
.input("src/model/my_model.onnx")
.out_dir("model/")
.run_from_script();
}
```
This generates Rust code and a `.bpk` weights file from your ONNX model during the build process.
### Step 3: Modify `mod.rs`
In your `src/model/mod.rs` file, include the generated code:
```rust, ignore
pub mod my_model {
include!(concat!(env!("OUT_DIR"), "/model/my_model.rs"));
}
```
### Step 4: Use the Imported Model
Now you can use the imported model in your code:
```rust, ignore
use burn::tensor;
use burn::backend::{Flex, flex::FlexDevice};
use model::my_model::Model;
fn main() {
let device = FlexDevice;
// Create model instance and load weights from target dir default device
let model: Model<Flex> = Model::default();
// Create input tensor (replace with your actual input)
let input = tensor::Tensor::<Flex, 4>::zeros([1, 3, 224, 224], &device);
// Perform inference
let output = model.forward(input);
println!("Model output: {:?}", output);
}
```
## Advanced Configuration
The `ModelGen` struct provides configuration options:
```rust, ignore
use burn_onnx::{ModelGen, LoadStrategy};
ModelGen::new()
.input("path/to/model.onnx")
.out_dir("model/")
.development(true) // Enable development mode for debugging
.load_strategy(LoadStrategy::Embedded) // Embed weights in the binary
.run_from_script();
```
- `input`: Path to the ONNX model file
- `out_dir`: Output directory for generated code and weights
- `development`: When enabled, generates additional debug files (`.onnx.txt`, `.graph.txt`)
- `load_strategy`: Controls which weight-loading constructors are generated on the `Model` struct
(see below)
Model weights are stored in Burnpack format (`.bpk`), which provides efficient serialization and
loading.
### Load Strategy
The `LoadStrategy` enum controls how the generated model loads its weights:
| Strategy | Generated constructors | `Default` impl | Use case |
|------------|------------------------------------------------|-----------------|-------------------------------------------|
| `File` | `from_file()`, `from_bytes()` | Yes | Standard desktop/server (default) |
| `Embedded` | `from_embedded()`, `from_bytes()` | Yes | Single binary, small models |
| `Bytes` | `from_bytes()` | No | WASM, embedded, custom loaders |
| `None` | (none) | No | Manual weight management |
The default strategy is `File`, which keeps weights in a separate `.bpk` file and generates a
`from_file()` constructor.
For WebAssembly or environments without filesystem access, use `LoadStrategy::Bytes`:
```rust, ignore
ModelGen::new()
.input("model.onnx")
.out_dir("model/")
.load_strategy(LoadStrategy::Bytes)
.run_from_script();
```
Then load weights at runtime from any byte source (e.g., a network fetch):
```rust, ignore
let model = Model::<Backend>::from_bytes(weight_bytes, &device);
```
## Loading and Using Models
You can load models in several ways, depending on the `LoadStrategy` used during code generation:
```rust, ignore
// Load from the output directory with default device (recommended for most use cases)
// This automatically loads weights from the .bpk file
// Available with LoadStrategy::File or LoadStrategy::Embedded
let model = Model::<Backend>::default();
// Create a new model instance with a specific device
// (initializes weights randomly; load weights via `load_from` afterward)
let model = Model::<Backend>::new(&device);
// Load from a specific .bpk file (LoadStrategy::File)
let model = Model::<Backend>::from_file("path/to/weights.bpk", &device);
// Load from in-memory bytes (LoadStrategy::File, Embedded, or Bytes)
let model = Model::<Backend>::from_bytes(weight_bytes, &device);
// Load from embedded weights (LoadStrategy::Embedded)
let model = Model::<Backend>::from_embedded(&device);
```
## Troubleshooting
Common issues and solutions:
1. **Unsupported ONNX operator**: Check the
[list of supported ONNX operators](https://github.com/tracel-ai/burn-onnx/blob/main/SUPPORTED-ONNX-OPS.md).
You may need to simplify your model or wait for support.
2. **Build errors**: Ensure your `burn-onnx` version matches your Burn version and verify the ONNX
file path in `build.rs`.
3. **Runtime errors**: Confirm that your input tensors match the expected shape and data type of
your model.
4. **Performance issues**: Consider using a more performant backend or optimizing your model
architecture.
5. **Viewing generated files**: Find the generated Rust code and weights in the `OUT_DIR` directory
(usually `target/debug/build/<project>/out`).
## Examples and Resources
For practical examples, check out the
[burn-onnx examples](https://github.com/tracel-ai/burn-onnx/tree/main/examples):
1. [ONNX Inference](https://github.com/tracel-ai/burn-onnx/tree/main/examples/onnx-inference) -
MNIST inference example
2. [Image Classification Web](https://github.com/tracel-ai/burn-onnx/tree/main/examples/image-classification-web) -
SqueezeNet running in the browser via WebAssembly
3. [Raspberry Pi Pico](https://github.com/tracel-ai/burn-onnx/tree/main/examples/raspberry-pi-pico) -
Embedded deployment example
These demonstrate real-world usage of ONNX import in Burn projects.
For contributors looking to add support for new ONNX operators:
- [Development Guide](https://github.com/tracel-ai/burn-onnx/blob/main/DEVELOPMENT-GUIDE.md) -
Step-by-step guide for implementing new operators
## Conclusion
Importing ONNX models into Burn combines the vast ecosystem of pre-trained models with Burn's
performance and Rust's safety features. Following this guide, you can seamlessly integrate ONNX
models into your Burn projects for inference, fine-tuning, or further development.
The `burn-onnx` crate is actively developed, with ongoing work to support more ONNX operators and
improve performance. Visit the [burn-onnx repository](https://github.com/tracel-ai/burn-onnx) for
updates and to contribute!
burn-book/src/overview.md
+37
View File
@@ -0,0 +1,37 @@
# Overview
Welcome to The Burn Book 👋
This book will help you get started with the Burn deep learning framework, whether you are an
advanced user or a beginner. We have crafted some sections for you:
- [Basic Workflow: From Training to Inference](./basic-workflow): We'll start with the fundamentals,
guiding you through the entire workflow, from training your models to deploying them for
inference. This section lays the groundwork for your Burn expertise.
- [Building Blocks](./building-blocks): Dive deeper into Burn's core components, understanding how
they fit together. This knowledge forms the basis for more advanced usage and customization.
- [Performance - Good Practices](./performance/good-practices/): Tips for writing models and
training code that make the most of hardware resources while avoiding common pitfalls that can
slow down execution.
- [Custom Training Loop](./custom-training-loop.md): Gain the power to customize your training
loops, fine-tuning your models to meet your specific requirements. This section empowers you to
harness Burn's flexibility to its fullest.
- [Saving & Loading Models](./saving-and-loading.md): Learn how to save and load your trained
models, including importing weights from PyTorch and SafeTensors formats.
- [ONNX Import](./onnx-import.md): Learn how to import ONNX models using the
[burn-onnx](https://github.com/tracel-ai/burn-onnx) crate.
- [Models & Pre-Trained Weights](./models-and-pretrained-weights.md): Get started quickly with
ready-to-use models and pre-trained weights.
- [Advanced](./advanced): Finally, venture into advanced topics, exploring Burn's capabilities at
their peak. This section caters to those who want to push the boundaries of what's possible with
Burn.
Throughout the book, we assume a basic understanding of deep learning concepts, but we may refer to
additional material when it seems appropriate.
burn-book/src/performance/README.md
+4
View File
@@ -0,0 +1,4 @@
# Performance
This section covers the key concepts you need to understand to get the most out of Burn and your
hardware.
burn-book/src/performance/distributed-computing.md
+4
View File
@@ -0,0 +1,4 @@
# Distributed Computing
Distributed computing support was introduced in Burn 0.19. Documentation and examples will be
available soon.
burn-book/src/performance/good-practices/README.md
+12
View File
@@ -0,0 +1,12 @@
# Performance - Best Practices
This section provides valuable insights into the performance characteristics of Burn and guides
users on how to effectively leverage them for optimal results.
It includes several sections, each offering relevant details. While understanding these concepts can
aid in model optimization, it’s always crucial to conduct benchmarks and profile models to
accurately assess performance improvements.
- [Asynchronous Execution](./asynchronous-execution.md)
- [Kernel Fusion](./kernel-fusion.md)
- [Kernel Selection](./kernel-selection.md)
burn-book/src/performance/good-practices/asynchronous-execution.md
+62
View File
@@ -0,0 +1,62 @@
# Asynchronous Execution
Most Burn backends execute tensor operations in an asynchronous manner. However, the async notation
is often not required for most tensor operations, privileging the simplicity of sync Rust.
There are only a few operations that trigger synchronization of the backend, and it is very
important to correctly handle those to optimize hardware utilization. Those operations are
`into_data`, `into_scalar`, and `sync`. Some tensor operations might call `into_data` underneath,
triggering a synchronization, like `to_device` for some backends.
There are several ways to minimize synchronization overhead, one of which is to batch sync
operations into a single transaction. Burn provides a high-level composable API to build
transactions, which will only trigger a single sync on the device.
For instance, it is often used when collecting metrics during training:
```rust
/// All of these variables are tensors.
let (output, loss, targets) = ..;
/// Now output, loss, and targets will be `TensorData` stored on the CPU.
let [output, loss, targets] = Transaction::default()
.register(output)
.register(loss)
.register(targets)
.execute()
.try_into()
.expect("Correct amount of tensor data");
```
Another way of optimizing reads and avoiding device stalls is to read the data on a different
thread. Under the hood, CubeCL-based backends assign different execution queues for different
threads, meaning that syncing a thread shouldn’t impact the throughput of another thread.
## Using Different Backends for Different Tasks
Tensor operations aren’t the only things that are asynchronous; dataset and dataloading are also
lazily executed. This allows for efficient data augmentation and sampling without having to cache
huge datasets on disk. However, this might reduce training throughput if data augmentation is
performed on the same device as the training itself. So, it is normally encouraged to use a
different device, maybe even a different backend, for that purpose. For optimal performance, also
avoid small allocations followed by a batching procedure. Even if it doesn’t break asynchronicity,
it can slow down performance.
```rust
/// Items is a vector of many tensors.
let items = ..;
let batch = Tensor::cat(items, 1);
```
Prefer doing the concatenation of tensors on the data augmentation device and not on the training
device.
```rust
/// Items is a vector of many tensors.
let items = ..;
let device_training = ..;
let axis_batch = 0;
let items = Tensor::cat(items, axis_batch);
let batch = Tensor::from_data(items.into_data(), device_training);
```
burn-book/src/performance/good-practices/kernel-fusion.md
+76
View File
@@ -0,0 +1,76 @@
# Kernel Fusion
An interesting property of async execution is that it allows performance optimizations like kernel
fusion. Coupled with CubeCL and its Just-In-Time compiler, Burn can serialize tensor operations into
a symbolic graph, then optimize it for improved efficiency.
Kernel fusion may reorder operations to reduce global memory reads, writes, and allocations. Being
aware of which operations can be fused is relevant, as it can be easy to break an execution graph.
The easiest way to optimize for fusion is to avoid keeping tensors alive for too long. When fusion
isn’t possible, all tensors that will be used later will trigger a global memory write. Fortunately,
Rust and Clippy are quite good at detecting unnecessary clones, but special care should still be
taken.
View operations can also interfere with fusion. They can be included in optimized graphs, but only
to a limited extent, and they reduce vectorization potential as we have fewer guarantees about
memory access patterns with transformed indices. So, it is good practice to group view operations
together before executing blocks of operations.
```rust
let tensor4 = tensor1.unsqueeze().matmul(tensor2) + tensor3.unsqueeze();
```
Could be improved with the following:
```rust
let tensor1 = tensor1.unsqueeze();
let tensor3 = tensor3.unsqueeze();
let tensor4 = tensor1.matmul(tensor2) + tensor3;
```
This reduces the necessary reordering and may reduce a global memory write or improve vectorization.
We might be able to detect these patterns in the future, but for now, it’s a good idea to order your
operations using this pattern. As a reminder, view operations typically only update tensor metadata
in most cases. These operations include `slice`, `slice_assign`, `select`, `gather`, `scatter`,
`reshape`, `swap_dims`, `transpose`, `unsqueeze`, etc.
With fusion enabled, it is often not necessary to write custom kernels, as you can rely on our
system to optimize most element-wise operations. However, most compute-bound kernels require many
tricks and deep knowledge of GPU memory architectures, where automatic compiler optimizations often
underperform compared to human-designed algorithms. This is why Burn’s approach to fusion is
centered around fuse-on-read and fuse-on-write. This means that complex compute-bound kernels that
change the shapes of tensors can fuse a block of element-wise operations when reading the input
tensor and when writing the output tensor. The implication is that multiple compute-bound operations
in a sequence can reduce fusion potential.
```rust
// This line might trigger 3 writes: tensor1, tensor2, and tensor3, if tensor1 and tensor2 are abstract tensors.
let tensor3 = tensor1.clone().sum_dim(tensor2.clone(), 2);
let tensor4 = tensor2.sum_dim(tensor3, 2);
let tensor5 = tensor4 + (tensor1 * tensor2);
```
```rust
let tmp = tensor1.clone() + tensor2.clone();
let tensor3 = tensor1.sum_dim(tensor2, 2);
let tensor4 = tensor2.sum_dim(tensor3, 2);
let tensor5 = tensor4 + tmp;
```
The lesson? Whenever possible, pass only the latest value to a compute operation. Don’t clone a
tensor before compute-bound operations, as it might trigger an additional write if that tensor isn’t
materialized from initial fusion.
It’s a bit complex, but the first code snippet is actually better if `tensor1` and `tensor2` are
concrete in global memory. This would be the case if `tensor1` and `tensor2` are model parameters,
so prefer this implementation style in such scenarios.
The second code snippet is preferred when `tensor1` and `tensor2` are virtual tensors, meaning they
were fused by earlier operations and require a global memory read to be accessed later. This happens
if those tensors are part of a signal in neural networks.
Reordering operations can help in such scenarios but will not create temporary values, making the
previous optimization harder. We might eventually automatically optimize these cases, but the
solution space is quite large, and it’s not a planned optimization. Profiling model blocks is always
a good idea to identify which code block is faster when faced with ambiguous situations.
burn-book/src/performance/good-practices/kernel-selection.md
+25
View File
@@ -0,0 +1,25 @@
# Kernel Selection
As mentioned earlier, complex compute-bound operations are highly non-trivial and require many
tricks for optimal performance. However, the way these tricks are applied varies depending on the
hardware and problem shapes. To select the best kernel, we use a search method with a highly
configurable autotune system that performs micro-benchmarks at runtime on the current hardware.
This may trigger a cold start, but the results of these benchmarks are cached on disk for subsequent
executions.
For deployment or training on spot instances, it’s a good idea to bundle the autotune cache with the
code to mitigate cold starts. Refer to the
[CubeCL configuration documentation](https://burn.dev/books/cubecl/advanced-usage/config.html) for
more details on fine-grained settings .
From the user’s point of view, kernel selection shouldn’t be a problem, but as usual, crafting
models with even shapes, multiples of 8, can significantly improve performance. Avoid creating
tensors with shapes that are multiples of 10, like `[1000, 1000]`, as these typically require bounds
checking and may limit vectorization.
Prefer shapes like `[1024, 1024]`, where dimensions are multiples of 32 or powers of 2, as these are
generally optimal. If you have no choice but to use a suboptimal shape, prefer handling it in a
single kernel, transforming it into an optimal shape. It’s better to have a slow neural network
layer followed by fast ones than to propagate unevenness and end up with smaller, but slower,
layers.
burn-book/src/performance/quantization.md
+133
View File
@@ -0,0 +1,133 @@
# Quantization
Quantization techniques perform computations and store tensors in lower precision data types like
8-bit integer instead of floating point precision. There are multiple approaches to quantize a deep
learning model categorized as:
- Post-training quantization (PTQ)
- Quantization aware training (QAT)
In post-training quantization, the model is trained in floating point precision and later converted
to the lower precision data type. There are two types of post-training quantization:
1. Static quantization: quantizes the weights and activations of the model. Quantizing the
activations statically requires data to be calibrated (i.e., recording the activation values to
compute the optimal quantization parameters with representative data).
1. Dynamic quantization: quantized the weights ahead of time (like static quantization) but the
activations are dynamically at runtime.
Sometimes post-training quantization is not able to achieve acceptable task accuracy. In general,
this is where quantization-aware training (QAT) can be used: during training, fake-quantization
modules are inserted in the forward and backward passes to simulate quantization effects, allowing
the model to learn representations that are more robust to reduced precision.
Burn does not currently support QAT. Only post-training quantization (PTQ) is implemented at this
time.
<div class="warning">
Quantization support in Burn is currently in active development.
It supports the following PTQ modes on some backends:
- Per-tensor and per-block quantization to 8-bit, 4-bit and 2-bit representations
No integer operations are currently supported, which means tensors are dequantized to perform the
operations in floating point precision.
</div>
## Module Quantization
Quantizing the weights of your model after training is quite simple. We have access to the weight
tensors and can collect their statistics, such as the min and max value when using
`MinMaxCalibration`, to compute the quantization parameters.
```rust , ignore
# use burn::module::Quantizer;
# use burn::tensor::quantization::{Calibration, QuantLevel, QuantParam, QuantScheme, QuantValue};
#
// Quantization config
let scheme = QuantScheme::default()
.with_level(QuantLevel::Block(32))
.with_value(QuantValue::Q4F)
.with_param(QuantParam::F16);
let mut quantizer = Quantizer {
calibration: Calibration::MinMax,
scheme,
};
// Quantize the weights
let model = model.quantize_weights(&mut quantizer);
```
### Calibration
Calibration is the step during quantization where the range of all floating-point tensors is
computed. This is pretty straightforward for weights since the actual range is known at
_quantization-time_ (weights are static), but activations require more attention.
To compute the quantization parameters, Burn supports the following `Calibration` methods.
| Method | Description |
| :------- | :------------------------------------------------------------------------------- |
| `MinMax` | Computes the quantization range mapping based on the running min and max values. |
### Quantization Scheme
A quantization scheme defines how an input is quantized, including the representation of quantized
values, storage format, granularity, and how the values are scaled.
```rust
let scheme = QuantScheme::default()
.with_mode(QuantMode::Symmetric) // Quantization mode
.with_level(QuantLevel::block([2, 16])) // Granularity (per-tensor or per-block)
.with_value(QuantValue::Q8S) // Data type of quantized values, independent of how they're stored
.with_store(QuantStore::Native) // Storage format for quantized values
.with_param(QuantParam::F16); // Precision for quantization parameters
```
#### Quantization Mode
| Mode | Description |
| :---------- | :------------------------------------------- |
| `Symmetric` | Values are scaled symmetrically around zero. |
#### Quantization Level
| Level | Description |
| :----------------------------- | :----------------------------------------------------------------------------------------------------------- |
| `Tensor` | A single quantization parameter set for the entire tensor. |
| `Block(block_size: BlockSize)` | Tensor divided into blocks (1D, 2D, or higher) defined by block_size, each with its own quantization params. |
#### Quantization Value
| Value | Bits | Description |
| :----- | :--: | :-------------------------------------------- |
| `Q8F` | 8 | 8-bit full-range quantization |
| `Q4F` | 4 | 4-bit full-range quantization |
| `Q2F` | 2 | 2-bit full-range quantization |
| `Q8S` | 8 | 8-bit symmetric quantization |
| `Q4S` | 4 | 4-bit symmetric quantization |
| `Q2S` | 2 | 2-bit symmetric quantization |
| `E5M2` | 8 | 8-bit floating-point (5 exponent, 2 mantissa) |
| `E4M3` | 8 | 8-bit floating-point (4 exponent, 3 mantissa) |
| `E2M1` | 4 | 4-bit floating-point (2 exponent, 1 mantissa) |
#### Quantization Store
| Store | Description |
| :------------------ | :------------------------------------------------------------------------------------------------------------------------------------------------ |
| `Native` | Each quantized value is stored directly in a native format, which doesn't require packing and unpacking. |
| `PackedNative(dim)` | Multiple quantized values packed into a 32-bit integer. Argument is the dimension the tensor is packed on, starting from the innermost dimension. |
| `PackedU32(dim)` | Multiple quantized values packed into a 32-bit integer. Argument is the dimension the tensor is packed on, starting from the innermost dimension. |
Native storage is not supported for sub-byte quantization values.
#### Quantization Parameters Precision
| Param | Description |
| :----- | :----------------------------- |
| `F32` | Full floating-point precision. |
| `F16` | Half-precision floating point. |
| `BF16` | Brain float 16-bit precision. |
burn-book/src/saving-and-loading.md
+468
View File
@@ -0,0 +1,468 @@
# Saving and Loading Models
Saving your trained machine learning model is quite easy. As mentioned in the
[Record](./building-blocks/record.md) section, a module's parameters are captured in a
`ModuleRecord` and serialized to the [burnpack](./building-blocks/record.md) format (`.bpk`).
```rust, ignore
use burn::store::ModuleRecord;
// Take a record of the model's parameters and save it to disk.
model
.into_record()
.save(model_path)
.expect("Should be able to save the model");
```
Note that the `.bpk` file extension is appended automatically when the path has none, so only the
file path and base name need to be provided.
Now that you have a trained model saved to your disk, you can just as easily load it back.
```rust, ignore
// Load the record from the burnpack file.
let record = ModuleRecord::load(model_path)
.expect("Should be able to load the model weights from the provided file");
// Apply the loaded weights to a model.
model = model.load_record(record);
```
The record is backend-independent: weights saved with one backend can be loaded on another. If you
need to convert precision, call `model.cast(dtype)` before taking the record, or use
`record.cast_to_module_dtype()` when loading.
## Initialization from Recorded Weights
The most straightforward way to load weights for a module is simply by using
[load_record](https://burn.dev/docs/burn/module/trait.Module.html#method.load_record). Note that
parameter initialization is lazy, therefore no actual tensor allocation and GPU/CPU kernels are
executed before the module is used. This means that you can use `init(device)` followed by
`load_record(record)` without any meaningful performance cost.
```rust, ignore
use burn::store::ModuleRecord;
// Create a dummy initialized model to save.
let device = Default::default();
let model = Model::init(&device);
// Save its parameters to a burnpack file.
model
.into_record()
.save(model_path)
.expect("Should be able to save the model");
```
Afterwards, the model can just as easily be loaded from the record saved on disk.
```rust, ignore
// Load the model record from the burnpack file.
let record = ModuleRecord::load(model_path)
.expect("Could not load model weights");
// Initialize a new model with the loaded record/weights.
let model = Model::init(&device).load_record(record);
```
For partial loading (only some parameters present in the record), use
`record.allow_partial(true)` before applying it, or `model.try_load_record(record)` for fallible
loading.
## Model Weight Store
While the `ModuleRecord` API above works well for basic saving and loading, the `burn-store` crate
adds memory efficiency and flexibility on top of the same burnpack format. It provides zero-copy
memory-mapped loading, cross-framework interoperability (PyTorch and SafeTensors), key remapping,
partial loading, and filtering.
### Supported Formats
| Format | Extension | Description |
| --------------- | -------------- | ----------------------------------------------------------------------------------------- |
| **Burnpack** | `.bpk` | Burn's native format with fast loading, zero-copy support, and training state persistence |
| **SafeTensors** | `.safetensors` | Industry-standard format from Hugging Face for secure tensor serialization |
| **PyTorch** | `.pt`, `.pth` | Direct loading of PyTorch model weights (read-only) |
### Saving a Model
```rust, ignore
use burn_store::{ModuleSnapshot, BurnpackStore};
// Save to Burnpack (recommended)
let mut store = BurnpackStore::from_file("model.bpk");
model.save_into(&mut store)?;
// Or save to SafeTensors
use burn_store::SafetensorsStore;
let mut store = SafetensorsStore::from_file("model.safetensors");
model.save_into(&mut store)?;
```
### Loading a Model
```rust, ignore
use burn_store::{ModuleSnapshot, BurnpackStore};
let device = Default::default();
let mut model = MyModel::init(&device);
// Load from Burnpack
let mut store = BurnpackStore::from_file("model.bpk");
model.load_from(&mut store)?;
```
### Loading from PyTorch
You can load weights directly from PyTorch `.pt` files:
```rust, ignore
use burn_store::{ModuleSnapshot, PytorchStore};
let mut model = MyModel::init(&device);
let mut store = PytorchStore::from_file("pytorch_model.pt");
model.load_from(&mut store)?;
```
#### Exporting from PyTorch
Save only the model weights (state_dict), not the entire model:
```python
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(2, 2, (2, 2))
self.conv2 = nn.Conv2d(2, 2, (2, 2), bias=False)
def forward(self, x):
return self.conv2(self.conv1(x))
model = Net()
torch.save(model.state_dict(), "model.pt") # Correct: save state_dict
# torch.save(model, "model.pt") # Wrong: saves entire model
```
#### Accessing Nested State Dicts
Some PyTorch checkpoints nest the state_dict under a key:
```rust, ignore
let mut store = PytorchStore::from_file("checkpoint.pt")
.with_top_level_key("state_dict");
model.load_from(&mut store)?;
```
### Loading from SafeTensors
For SafeTensors files exported from PyTorch, use the adapter for proper weight transformation:
```rust, ignore
use burn_store::{ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore};
let mut model = MyModel::init(&device);
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_from_adapter(PyTorchToBurnAdapter);
model.load_from(&mut store)?;
```
For SafeTensors files created by Burn, no adapter is needed:
```rust, ignore
let mut store = SafetensorsStore::from_file("model.safetensors");
model.load_from(&mut store)?;
```
#### Exporting from PyTorch to SafeTensors
```python
from safetensors.torch import save_file
model = Net()
save_file(model.state_dict(), "model.safetensors")
```
### Saving for PyTorch Compatibility
Use the adapter when saving for PyTorch consumption:
```rust, ignore
use burn_store::{BurnToPyTorchAdapter, SafetensorsStore};
let mut store = SafetensorsStore::from_file("for_pytorch.safetensors")
.with_to_adapter(BurnToPyTorchAdapter)
.skip_enum_variants(true);
model.save_into(&mut store)?;
```
### Handling Load Results
The `load_from` method returns detailed information about the loading process.
> **Note:** Inspecting `result.missing`, `result.errors`, etc. requires the store to be configured with
> [`.allow_partial(true)`](#partial-loading). Without it, a missing tensor causes a hard `Err`
> before you ever receive an `ApplyResult`.
```rust, ignore
// Use .allow_partial(true) to get an ApplyResult with structured info
let mut store = PytorchStore::from_file("pretrained.pt")
.allow_partial(true);
let result = model.load_from(&mut store)?;
// Print a formatted summary with suggestions
println!("{}", result);
// Or inspect individual fields
println!("Applied: {} tensors", result.applied.len());
println!("Missing: {:?}", result.missing);
println!("Errors: {:?}", result.errors);
if result.is_success() {
println!("All tensors loaded successfully");
}
```
If you don't use `.allow_partial(true)`, use normal error handling:
```rust, ignore
match model.load_from(&mut store) {
Ok(result) => println!("Loaded successfully: {} tensors", result.applied.len()),
Err(e) => eprintln!("Failed to load: {e}"),
}
```
### Adding Metadata
Burnpack and SafeTensors support custom metadata:
```rust, ignore
let mut store = BurnpackStore::from_file("model.bpk")
.metadata("version", "1.0")
.metadata("description", "My trained model")
.metadata("epochs", "100");
model.save_into(&mut store)?;
```
### Advanced Features
#### Key Remapping
Remap parameter names using regex patterns when model structures don't match:
```rust, ignore
let mut store = PytorchStore::from_file("model.pt")
// Remove prefix: "model.conv1.weight" -> "conv1.weight"
.with_key_remapping(r"^model\.", "")
// Rename: "layer1" -> "encoder.layer1"
.with_key_remapping(r"^layer", "encoder.layer");
model.load_from(&mut store)?;
```
For complex remapping:
```rust, ignore
use burn_store::KeyRemapper;
let remapper = KeyRemapper::new()
.add_pattern(r"^transformer\.h\.(\d+)\.", "transformer.layer$1.")?
.add_pattern(r"\.attn\.", ".attention.")?;
let mut store = SafetensorsStore::from_file("model.safetensors")
.remap(remapper);
```
#### Partial Loading
Load weights even when some tensors are missing:
```rust, ignore
let mut store = PytorchStore::from_file("pretrained.pt")
.allow_partial(true);
let result = model.load_from(&mut store)?;
println!("Missing (initialized randomly): {:?}", result.missing);
```
#### Filtering Tensors
Load or save only specific layers:
```rust, ignore
// Load only encoder layers
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_regex(r"^encoder\..*")
.allow_partial(true);
// Save only encoder layers
let mut store = SafetensorsStore::from_file("encoder.safetensors")
.with_regex(r"^encoder\..*");
model.save_into(&mut store)?;
// Multiple patterns (OR logic)
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_regex(r"^encoder\..*") // encoder tensors
.with_regex(r".*\.bias$") // OR any bias tensors
.with_full_path("decoder.scale"); // OR specific tensor
```
#### Non-Contiguous Layer Indices
PyTorch `nn.Sequential` with mixed layers creates non-contiguous indices. `PytorchStore`
automatically remaps these:
```
PyTorch: fc.0.weight, fc.2.weight, fc.4.weight (gaps from ReLU layers)
Burn: fc.0.weight, fc.1.weight, fc.2.weight (contiguous)
```
This is enabled by default. Disable if needed:
```rust, ignore
let mut store = PytorchStore::from_file("model.pt")
.map_indices_contiguous(false);
```
#### Zero-Copy Loading
For embedded models or large files, use zero-copy loading to avoid memory copies:
```rust, ignore
// Embedded model (compile-time)
static MODEL_DATA: &[u8] = include_bytes!("model.bpk");
let mut store = BurnpackStore::from_static(MODEL_DATA);
model.load_from(&mut store)?;
// Large file (memory-mapped)
let mut store = BurnpackStore::from_file("large_model.bpk")
.zero_copy(true);
model.load_from(&mut store)?;
```
#### Half-Precision Storage
Save models at half precision (F16) to reduce file size by ~50%, then load back at full precision:
```rust, ignore
use burn_store::{ModuleSnapshot, BurnpackStore, HalfPrecisionAdapter};
let adapter = HalfPrecisionAdapter::new();
// Save: F32 -> F16 (same adapter for both directions)
let mut store = BurnpackStore::from_file("model_f16.bpk")
.with_to_adapter(adapter.clone());
model.save_into(&mut store)?;
// Load: F16 -> F32
let mut store = BurnpackStore::from_file("model_f16.bpk")
.with_from_adapter(adapter);
model.load_from(&mut store)?;
```
By default, weights in Linear, Embedding, Conv\*, LayerNorm, GroupNorm, InstanceNorm, RmsNorm, and
PRelu modules are converted. BatchNorm is excluded because its running variance can underflow in
F16. Customize with `with_module()` and `without_module()`:
```rust, ignore
// Keep LayerNorm at full precision
let adapter = HalfPrecisionAdapter::new()
.without_module("LayerNorm");
// Add a custom module to the conversion set
let adapter = HalfPrecisionAdapter::new()
.with_module("CustomLayer");
```
#### Direct Tensor Access
Inspect tensors without loading into a model:
```rust, ignore
use burn_store::ModuleStore;
let mut store = PytorchStore::from_file("model.pt");
// List all tensor names
let names = store.keys()?;
// Get specific tensor
if let Some(snapshot) = store.get_snapshot("encoder.layer0.weight")? {
println!("Shape: {:?}, DType: {:?}", snapshot.shape, snapshot.dtype);
}
```
#### Model Surgery
Transfer weights between models:
```rust, ignore
use burn_store::{ModuleSnapshot, PathFilter};
// Transfer all weights
let snapshots = model1.collect(None, None, false);
model2.apply(snapshots, None, None, false);
// Transfer only encoder weights
let filter = PathFilter::new().with_regex(r"^encoder\..*");
let snapshots = model1.collect(Some(filter.clone()), None, false);
model2.apply(snapshots, Some(filter), None, false);
```
### API Reference
#### Builder Methods
| Category | Method | Description |
| ------------- | ------------------------------ | ---------------------------- |
| **Filtering** | `with_regex(pattern)` | Filter by regex pattern |
| | `with_full_path(path)` | Include specific tensor |
| | `with_predicate(fn)` | Custom filter logic |
| **Remapping** | `with_key_remapping(from, to)` | Regex-based renaming |
| | `remap(KeyRemapper)` | Complex remapping rules |
| **Adapters** | `with_from_adapter(adapter)` | Loading transformations |
| | `with_to_adapter(adapter)` | Saving transformations |
| | `HalfPrecisionAdapter::new()` | F32/F16 mixed-precision |
| **Config** | `allow_partial(bool)` | Continue on missing tensors |
| | `with_top_level_key(key)` | Access nested dict (PyTorch) |
| | `skip_enum_variants(bool)` | Skip enum variants in paths |
| | `map_indices_contiguous(bool)` | Remap non-contiguous indices |
| | `metadata(key, value)` | Add custom metadata |
| | `zero_copy(bool)` | Enable zero-copy loading |
#### Direct Access Methods
| Method | Description |
| --------------------- | -------------------------------- |
| `keys()` | Get ordered list of tensor names |
| `get_all_snapshots()` | Get all tensors as BTreeMap |
| `get_snapshot(name)` | Get specific tensor by name |
### Troubleshooting
#### Common Issues
1. **"Missing source values" error**: You saved the entire PyTorch model instead of the state_dict.
Re-export with `torch.save(model.state_dict(), "model.pt")`.
2. **Shape mismatch**: Your Burn model doesn't match the source architecture. Verify layer
configurations (channels, kernel sizes, bias settings).
3. **Key not found**: Parameter names don't match. Use `with_key_remapping()` or inspect keys:
```rust, ignore
let store = PytorchStore::from_file("model.pt");
println!("Available keys: {:?}", store.keys()?);
```
#### Inspecting Files
Use [Netron](https://github.com/lutzroeder/netron) to visualize `.pt` and `.safetensors` files.
For Burnpack files:
```bash
cargo run --example burnpack_inspect model.bpk
```
codecov.yml
+13
View File
@@ -0,0 +1,13 @@
coverage:
status:
project:
default:
# https://docs.codecov.com/docs/commit-status#informational
informational: true
target: 80%
patch:
default:
informational: true
target: 80%
github_checks:
annotations: false
contributor-book/.gitignore
+18
View File
@@ -0,0 +1,18 @@
target
# MacOS temp file
.DS_Store
book-test
guide/book
.vscode
tests/burn-book/book/
book/
# Ignore Jetbrains specific files.
.idea/
# Ignore Vim temporary and swap files.
*.sw?
*~
contributor-book/.prettierrc.json
+4
View File
@@ -0,0 +1,4 @@
{
"printWidth": 100,
"proseWrap": "always"
}
contributor-book/LICENSE-APACHE
Symlink
+1
View File
@@ -0,0 +1 @@
../LICENSE-APACHE
contributor-book/LICENSE-MIT
Symlink
+1
View File
@@ -0,0 +1 @@
../LICENSE-MIT
contributor-book/book.toml
+16
View File
@@ -0,0 +1,16 @@
[book]
authors = [
"Wouter Doppenberg",
"Nathaniel Simard",
"Louis Fortier-Dubois",
"Dilshod Tadjibaev",
"Guillaume Lagrange",
"Joshua Ferguson",
"The Burn Community",
]
language = "en"
src = "src"
title = "The Burn Contributor Book 🔥"
[output.html]
mathjax-support = true
contributor-book/src/SUMMARY.md
+16
View File
@@ -0,0 +1,16 @@
- [Overview](./overview.md)
- [How to Read This Book](./how-to-read-this-book.md)
- [Getting Started](./getting-started/README.md)
- [Setting Up The Environment](./getting-started/setting-up-the-environment.md)
- [Configuring Your Editor (Optional)](./getting-started/configuring-your-editor.md)
- [Testing](./getting-started/testing.md)
- [Architecture Overview](./project-architecture/README.md)
- [Modules](./project-architecture/module.md)
- [Serialization](./project-architecture/serialization.md)
- [Tensor](./project-architecture/tensor.md)
- [Backend](./project-architecture/backend.md)
- [Guides for Contributors](./guides/README.md)
- [Adding a New Operation to Burn](./guides/adding-a-new-operation-to-burn.md)
- [Submitting Examples to Burn](./guides/submitting-examples.md)
- [Frequently Encountered Issues](./frequently-encountered-issues/README.md)
- [Issues Related To Adding Operators](./frequently-encountered-issues/issues-while-adding-ops.md)
contributor-book/src/frequently-encountered-issues/README.md
+5
View File
@@ -0,0 +1,5 @@
# Frequently Encountered Issues
This is a collection of issues people have encountered and asked about on the
[Discord server](https://discord.gg/uPEBbYYDB6). This section is separated from the guides since it
can involve lots of details that are only relevant to a small subset of contributors.
contributor-book/src/frequently-encountered-issues/issues-while-adding-ops.md
+19
View File
@@ -0,0 +1,19 @@
# Issues encountered while adding ops
Below are some of the issues that were encountered while adding ops to the project. If you encounter
an issue while adding an op that isn't listed here, and it's not obvious how to fix it, you can add
it to this list or reach out on the [Discord server](https://discord.gg/uPEBbYYDB6) if you need
help.
## Off by .000001 errors
```sh
---- fusion::base::tests::maxmin::tests::test_mean_dim_2d stdout ---- thread 'fusion::base::tests::maxmin::tests::test_mean_dim_2d' panicked at burn-wgpu/src/fusion/base.rs:185:5: assertion `left == right` failed left: Data { value: [1.0, 4.0], shape: Shape { dims: [2, 1] } } right: Data { value: [0.99999994, 3.9999998], shape: Shape { dims: [2, 1] } } ----
tests::maxmin::tests::test_mean_dim_2d stdout ---- thread 'tests::maxmin::tests::test_mean_dim_2d' panicked at burn-wgpu/src/lib.rs:49:5: assertion `left == right` failed left: Data { value: [1.0, 4.0], shape: Shape { dims: [2, 1] } } right: Data { value: [0.99999994, 3.9999998], shape: Shape { dims: [2, 1] } }
```
If you encounter this, swap out the `assert_eq!` in the failing test for
`tensor1.to_data().assert_approx_eq` with `3` as the second argument. The second arguments specifies
the level of precision: `3` is equivalent to a less than 10<sup>-3</sup> (0.001) difference between
the elements of the two tensors.
contributor-book/src/getting-started/README.md
+6
View File
@@ -0,0 +1,6 @@
# Getting Started
This section is for setting up the environment and how to do basic development tasks such as running
tests and checking your code before committing. If you need help with the process or run into
issues, feel free to ask on the [Discord server](https://discord.gg/uPEBbYYDB6) in the Development
channels.
contributor-book/src/getting-started/configuring-your-editor.md
+36
View File
@@ -0,0 +1,36 @@
# Configuring your editor
These steps are not required, and most of this isn't specific to Burn, but it's definitely helpful
if you haven't already done it.
## VSCode
Install the following extensions:
- [rust-lang.rust-analyzer](https://marketplace.visualstudio.com/items?itemName=rust-lang.rust-analyzer)
for Rust syntax and semantic analysis
- [tamasfe.even-better-toml](https://marketplace.visualstudio.com/items?itemName=tamasfe.even-better-toml)
for TOML syntax and semantic analysis
- [fill-labs.dependi](https://marketplace.visualstudio.com/items?itemName=fill-labs.dependi) for
managing dependencies
- [vadimcn.vscode-lldb](https://marketplace.visualstudio.com/items?itemName=vadimcn.vscode-lldb) for
debugging
### Setting up the Debugger
To use the debugger, follow these steps:
1. Open `Command Palette` with `Ctrl+Shift+P` or `F1` and type
`LLDB: Generate Launch Configurations from Cargo.toml` then select it, this will generate a file
that should be saved as `.vscode/launch.json`.
2. Select the configuration from the "run and debug" side panel, then select the target from the list.
Since this repo has `debug = 0` in the root `Cargo.toml` to speed up compilation, you need replace it with `debug = true` in the root `Cargo.toml` when using a debugger and breakpoints with `launch.json` settings.
3. Now you can enable breakpoints on code through IDE then start debugging the library/binary you
want, like in the following example:
![debug-options](debug-options-vscode.png)
If you're creating a new library or binary, keep in mind to repeat step 1 to always keep a fresh
list of targets.
## Have another editor? Open a PR!
contributor-book/src/getting-started/debug-options-vscode.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 189 KiB

contributor-book/src/getting-started/setting-up-the-environment.md
+54
View File
@@ -0,0 +1,54 @@
# Setting up the environment
Depending on what part of the project you plan on contributing to, there are a couple of tools to
install and commands to be familiar with. This section should be up to date with current project
practices (as of 2024-04-15).
## General
There are a few commands you will want to run prior to any commit for a non-draft PR:
1. `cargo fmt --all` will run `rustfmt` on all files in the project.
2. `cargo clippy --fix` will run [Clippy](https://github.com/rust-lang/rust-clippy) and fix any
coding issues it can. Clippy necessitates to be in a clean Git state, but this can be
circumvented by adding the `--allow-dirty` flag.
3. `cargo run-checks` is a command used to test the project. It is required to run successfully
prior to merging a PR. Fair warning, running these tests can take a while[^linux_mem_note].
> Want more detailed macro error diagnostics? This is especially useful for debugging tensor-related tests:
>
> ```bash
> RUSTC_BOOTSTRAP=1 RUSTFLAGS="-Zmacro-backtrace" cargo run-checks
> ```
## Updating the burn semver version
If for some reason you need to bump for the next version (though that should probably be left to the
maintainers), edit the semantic version number in `burn/Cargo.toml`, and then run `cargo update` to
update the lock file.
## Contributing to either the Burn Book or Contributor Book
Both the Burn Book and the Contributor Book are built with mdbook. To open the book locally, run
`mdbook serve <path/to/book>` or `cargo xtask books {burn|contributor} open` which will install and
use mdbook automatically.
Alternatively, if you want to install mdbook directly, run the following command[^update_note]:
```bash
cargo install mdbook
```
Also instead of running `cargo run-checks`, you can run `cargo xtask check typos` to only check
for misspellings. This will install [typo](https://crates.io/crates/typos-cli), and if any are
encountered you should be able to run `typo -w /path/to/book` to fix them.
[^linux_mem_note]:
If your system is running into issues with memory and you are on linux, you may want to switch
to a [virtual console](https://wiki.archlinux.org/title/Linux_console#Virtual_consoles) to run
the tests. To do this, press `ctrl+alt+f3` to switch to a virtual console (and log in), and
either `ctrl+alt+f1` or `ctrl+alt+f2` to switch back to your graphical session.
[^update_note]:
You might also want to install [cargo-update](https://github.com/nabijaczleweli/cargo-update) to
easily keep your tools up to date, though it is in no way required.
contributor-book/src/getting-started/testing.md
+38
View File
@@ -0,0 +1,38 @@
# Testing
## Test for Tensor Operations
Test for tensor operations (generally of the form: given this input, expect it match or approximate
this output) are defined only in
[`crates/burn-tensor/src/test/ops`](https://github.com/tracel-ai/burn/tree/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-tensor/src/tests/ops)
and not in the backends (with the exception of `burn-autodiff`). The tensor operation tests are
added to the `testgen_all` macro rule in
[`crates/burn-tensor/src/tests/mod.rs`](https://github.com/tracel-ai/burn/blob/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-tensor/src/tests/mod.rs).
This is then propagated to the existing backends without any additional work.
### Test for Autodiff
Tests for autodiff go under
[burn-autodiff/src/tests](https://github.com/tracel-ai/burn/tree/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-autodiff/src/tests)
and should verify backward pass correctness. For binary tensor operations, both the left and right
sides need to be verified.
Here's an easy way to define tests for a new operation's backward pass:
1. Use small tensors with simple values.
2. Pop open a terminal, launch `ipython` and import `numpy` then do the calculations by hand. You
can also use [Google Colab](https://colab.google/) so you don't have to install the packages on
your system.
3. Compare the actual outputs to the expected output for left-hand side, right-hand side.
For float tensors, it is advised to use
`actual_output_tensor.into_data().assert_approx_eq::<FloatElem<TestBackend>>(&expected_tensor_data, Tolerance::default())`
instead of `assert_eq!(...` due to occasional hiccups with floating point calculations. Other
assertions should also always use `FloatElem<TestBackend>`, and use `.elem()` to convert any
literals. Backends are tested for multiple precisions, and hardcoding to a fixed type causes tests
to fail with alternate floating point precisions. For convenience, it might be worth aliasing the
type like `type FT = FloatElem<TestBackend>;`.
For integers, tests should use `IntElem<TestBackend>`, and exit the test if the test values are
unrepresentable (above `max_value`, below `min_value`). A minimum range of `[0..127]` (`i8`) can be
assumed.
contributor-book/src/guides/README.md
+3
View File
@@ -0,0 +1,3 @@
# Guides for Contributors
The following guides are meant to help contributors accomplish specific tasks, such as adding new operations to Burn.

Some files were not shown because too many files have changed in this diff Show More