Создадим с нуля небольшой фреймворк машинного обучения на Rust.
Цели
1. Выяснить, заметен ли рост скорости при переходе с Python и PyTorch на Rust и LibTorch, серверную библиотеку PyTorch на C++, особенно в процессе обучения модели. ML-модели становятся крупнее, для их обучения требуется больше вычислительных возможностей, для обычного человека порой нереальных. Один из способов уменьшить рост аппаратных требований — понять, как сделать алгоритмы вычислительно эффективнее. Python в PyTorch — это лишь слой поверх LibTorch. Вопрос в том, стоит ли менять его на Rust. Планировалось использовать крейт Tch-rs для доступа к тензорам и функционалу Autograd DLL-библиотеки LibTorch как «калькулятору градиентов», а затем разработать с нуля на Rust остальное.
2. Сделать код достаточно простым для четкого понимания всех вычислений линейной алгебры и с возможностью легко его расширить при необходимости.
3. Во фреймворке ML-модели должны определяться, насколько это возможно, по аналогичной структуре стандартных Python/PyTorch.
4. Поизучать Rust и не скучать.
Но статья посвящена скорее преимуществам применения Rust в машинном обучении.
• Переходим сразу к конечному результату — вот как маленьким фреймворком создаются нейросетевые модели.
Листинг 1. Определение нейросетевой модели
struct MyModel {
l1: Linear,
l2: Linear,
}
impl MyModel {
fn new (mem: &mut Memory) -> MyModel {
let l1 = Linear::new(mem, 784, 128);
let l2 = Linear::new(mem, 128, 10);
Self {
l1: l1,
l2: l2,
}
}
}
impl Compute for MyModel {
fn forward (&self, mem: &Memory, input: &Tensor) -> Tensor {
let mut o = self.l1.forward(mem, input);
o = o.relu();
o = self.l2.forward(mem, &o);
o
}
}• Затем модель инстанцируется и обучается.
Листинг 2. Инстанцирование и обучение нейросетевой модели
fn main() {
let (x, y) = load_mnist();
let mut m = Memory::new();
let mymodel = MyModel::new(&mut m);
train(&mut m, &x, &y, &mymodel, 100, 128, cross_entropy, 0.3);
let out = mymodel.forward(&m, &x);
println!("Training Accuracy: {}", accuracy(&y, &out));
}Для пользователей PyTorch это интуитивно понятная аналогия определения и обучения нейросети на Python. В примере выше показана модель нейросети, используемая затем для классификации. Модель применяется к набору данных Mnist тестов производительности для сравнения двух версий модели: Rust и Python.
• В первом блоке кода создается структура MyModel с двумя слоями типа Linear.
• Второй блок — ее реализация, где определяется ассоциированная функция new, которой инициализируются два слоя и возвращается новый экземпляр структуры.
• В третьем блоке реализуется типаж Compute для MyModel, им определяется метод forward. Затем в функции main загружается набор данных Mnist, инициализируется память, инстанцируется MyModel, а после она обучается в течение 100 эпох с размером пакета 128, потерями перекрестной энтропии и скоростью обучения 0,3.
Очень даже понятно: это то, что потребуется для создания и обучения новых моделей на Rust с помощью маленького фреймворка. Теперь копнем поглубже и разберемся, как это все возможно.
Если вы привыкли создавать ML-модели в PyTorch, то наверняка, глядя на код выше, зададитесь вопросом: «Зачем здесь ссылка на Memory?». Объясним ниже.
Читать