Новый фреймворк глубокого обучения, полностью построенный на языке Rust, который призван обеспечить баланс между гибкостью, производительностью и простотой использования для исследователей, инженеров ML и разработчиков.

Что такое Rust Burn?

Rust Burn – это новый фреймворк глубокого обучения, полностью написанный на языке программирования Rust. Мотивация создания нового фреймворка вместо использования уже существующих, таких как PyTorch или TensorFlow, заключается в создании универсального фреймворка, ориентированного на различных пользователей, включая исследователей, инженеров машинного обучения и разработчиков низкоуровневого программного обеспечения.

Ключевыми принципами проектирования Rust Burn являются гибкость, производительность и простота использования.

Гибкость обусловлена возможностью быстрой реализации передовых исследовательских идей и проведения экспериментов.

Производительность достигается за счет оптимизаций, таких как использование специфических аппаратных возможностей, например, тензорных ядер на графических процессорах Nvidia.

Простота использования обусловлена упрощением рабочего процесса обучения, развертывания и запуска моделей в производство.

Основные характеристики:

Гибкий и динамический вычислительный граф
Потокобезопасные структуры данных
Интуитивно понятные абстракции для упрощения процесса разработки
Потрясающе высокая производительность при обучении и выводах
Поддержка нескольких реализаций бэкэнда для CPU и GPU
Полная поддержка протоколирования, метрики и контрольных точек во время обучения
Небольшое, но активное сообщество разработчиков

Начало работы

Установка ржавчины

Burn – это мощный фреймворк глубокого обучения, основанный на языке программирования Rust. Для работы с ним требуется базовое понимание Rust, но, освоив его, вы сможете воспользоваться всеми возможностями Burn.

Для установки можно воспользоваться официальным руководством. Вы также можете ознакомиться с руководством GeeksforGeeks по установке Rust на Windows и Linux со скриншотами.

Изображение из Install Rust

Установка Burn

Для использования Rust Burn необходимо, чтобы Rust был установлен на вашей системе. После того как Rust правильно установлен, можно создать новое приложение Rust с помощью cargo, менеджера пакетов Rust.

Выполните следующую команду в текущем каталоге:

cargo new new_burn_app

Перейдите в этот новый каталог:

cd new_burn_app

Затем добавьте Burn в качестве зависимости, а также функцию бэкенда WGPU, которая позволяет выполнять операции на GPU:

cargo add burn --features wgpu

В завершение скомпилируйте проект для установки Burn:

cargo build

В результате будет установлен фреймворк Burn вместе с бэкендом WGPU. WGPU позволяет Burn выполнять низкоуровневые операции на GPU.

Пример кода

Дополнение Element Wise

Для выполнения следующего кода необходимо открыть и заменить содержимое в файле src/main.rs:

use burn::tensor::Tensor;
use burn::backend::WgpuBackend;

// Type alias for the backend to use.
type Backend = WgpuBackend;

fn main() {
    // 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::::from_data([[2., 3.], [4., 5.]]);
    let tensor_2 = Tensor::::ones_like(&tensor_1);

    // Print the element-wise addition (done with the WGPU backend) of the two tensors.
    println!("{}", tensor_1 + tensor_2);
}

В функции main мы создали два тензора с бэкендом WGPU и выполнили сложение.

Для выполнения кода необходимо запустить в терминале команду cargo run.

Выход:

Теперь вы должны иметь возможность просмотреть результат добавления.

Tensor {
  data: [[3.0, 4.0], [5.0, 6.0]],
  shape:  [2, 2],
  device:  BestAvailable,
  backend:  "wgpu",
  kind:  "Float",
  dtype:  "f32",
}

Примечание: приведенный ниже код является примером из книги Burn Book: Начало работы.

Позиционный модуль Feed Forward

Приведем пример того, как просто использовать фреймворк. С помощью этого фрагмента кода мы объявляем позиционный модуль прямой передачи и его прямой проход.

use burn::nn;
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 PositionWiseFeedForward<B> {
    pub fn forward(&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)
    }
}

Приведенный выше код взят из репозитория GitHub.

Примеры проектов

Чтобы ознакомиться с другими примерами и запустить их, клонируйте репозиторий https://github.com/burn-rs/burn и запустите приведенные ниже проекты:

MNIST: Обучение модели на CPU или GPU с использованием различных бэкендов.
MNIST Inference Web: Вывод модели в браузере.
Классификация текста: Обучение кодера-трансформера с нуля на GPU.
Генерация текста: Построение и обучение авторегрессионного трансформатора с нуля на GPU.

Предварительно обученные модели

Для создания приложения искусственного интеллекта можно использовать следующие предварительно обученные модели и произвести их тонкую настройку на основе вашего набора данных.

SqueezeNet: squeezenet-burn
Llama 2: Gadersd/llama2-burn
Whisper: Gadersd/whisper-burn
Stable Diffusion v1.4: Gadersd/stable-diffusion-burn

Заключение

Rust Burn представляет собой новую интересную возможность в области фреймворков глубокого обучения. Если вы уже являетесь разработчиком Rust, вы можете использовать скорость, безопасность и параллелизм Rust, чтобы расширить границы возможного в исследованиях и производстве глубокого обучения. Burn стремится найти оптимальный компромисс между гибкостью, производительностью и удобством использования, чтобы создать уникальный универсальный фреймворк, подходящий для самых разных случаев применения.

Пока Burn находится на ранней стадии развития, он демонстрирует многообещающие возможности по устранению “болевых точек” существующих фреймворков и удовлетворению потребностей различных специалистов в этой области. По мере становления фреймворка и роста сообщества вокруг него, он может стать готовым к производству фреймворком наравне с уже существующими вариантами. Его новый дизайн и выбор языка открывают новые возможности для сообщества глубокого обучения.