Рост производительности машинного обучения с Rust. Часть 2

В первой части мы начали с нуля разработку ML-фреймворка на Rust. Основная цель  —  оценить рост скорости обучения модели на Rust и PyTorch по сравнению с Python. Для сетей прямой связи результаты очень позитивные.

Продолжим со сверточными нейросетями CNN, их определением и обучением. Чтобы иметь доступ к тензорной линейной алгебре и функционалу Autograd, снова задействуем крейт Tch-rs как обертку LibTorch, библиотеки PyTorch на C++. Остальное разрабатывается с нуля.

Код для обеих частей доступен на Github. Конечный результат  —  определение сверточных нейросетей CNN на Rust.

Листинг 1. Определение модели CNN

struct MyModel {
    l1: Conv2d,
    l2: Conv2d,
    l3: Linear,
    l4: Linear,
}

impl MyModel {
    fn new (mem: &mut Memory) -> MyModel {
        let l1 = Conv2d::new(mem, 5, 1, 10, 1);
        let l2 = Conv2d::new(mem, 5, 10, 20, 1);
        let l3 = Linear::new(mem, 320, 64);
        let l4 = Linear::new(mem, 64, 10);
        Self {
            l1: l1,
            l2: l2,
            l3: l3,
            l4: l4,
        }
    }
}

impl Compute for MyModel {
    fn forward (&self,  mem: &Memory, input: &Tensor) -> Tensor {
        let mut o = self.l1.forward(mem, &input);
        o = o.max_pool2d_default(2);
        o = self.l2.forward(mem, &o);
        o = o.max_pool2d_default(2);
        o = o.flat_view();
        o = self.l3.forward(mem, &o);
        o = o.relu();
        o = self.l4.forward(mem, &o);
        o
    }
}

… затем инстанцирование и обучение:

Листинг 2. Обучение модели CNN

fn main() {
    let (mut x, y) = load_mnist();
    x = x / 250.0;
    x = x.view([-1, 1, 28, 28]);
 
    let mut m = Memory::new();
    let mymodel = MyModel::new(&mut m);    
    train(&mut m, &x, &y, &mymodel, 20, 512, cross_entropy, 0.0001);
    let out = mymodel.forward(&m, &x);
    println!("Accuracy: {}", accuracy(&y, &out));
}

Чтобы листинг 1 был интуитивно понятен пользователям Python и PyTorch, определение модели максимально приближено к соответственному эквиваленту.

В структуре MyModel теперь добавляются слои Conv2D, которые затем инициируются в ассоциированной функции new.

В реализации типажа Compute определяется функция forward, которой принимаются входные данные через все слои, включая промежуточную функцию MaxPooling.

В функции main листинга 2, аналоге main из первой части, модель обучается и применяется к набору данных Mnist.

Рассмотрим подробнее такое определение и обучение CNN, это новые добавления к начатому в первой части фреймворку.

Поэкспериментируем с ядрами

Уникальная характеристика сверточных сетей  —  как минимум в одном слое вместо общего перемножения матриц применяется свертка. Цель операции свертки  —  с помощью ядра извлечь из входного изображения нужные признаки. Ядро  —  это матрица, перемещаемая по подразделам входного изображения и перемножаемая с ними так, что входные данные преобразуются в выходные определенным образом:

В двумерном случае в качестве входных данных используются двумерное изображение I и, как правило, двумерное ядро K, в итоге получаются такие вычисления свертки:

По уравнению 1 ясно, что наивный алгоритм применения свертки довольно затратен вычислительно из-за большого числа контуров и матричных умножений. К тому же эти вычисления многократно повторяются для каждого сверточного слоя сети и каждого примера/пакета обучения. Поэтому, прежде чем расширить для работы с CNN библиотеку из части 1, находим эффективный способ вычисления сверток.

Это задача очень хорошо изучена. Из различных вариантов, в том числе версий исключительно на Rust с промежуточными преобразованиями данных из тензоров PyTorch, решено использовать функцию свертки из LibTorch на C++. Для нее создается небольшая программа, в которой принимаемое цветное изображение преобразуется в полутоновое, а затем для выделения границ применяются известные ядра.

Сначала в чате Bing от Microsoft генерируем изображение. Применяем функцию свертки с гауссовым и лапласовым ядрами:

Ядра применены с методом conv2d, который берется из LibTorch на C++ через Tch-rs в таком виде:

Листинг 3. Метод «conv2d», берется из LibTorch через «Tch-rs»

pub fn conv2d<T: Borrow<Tensor>>(
    &self,
    weight: &Tensor,
    bias: Option<T>,
    stride: impl IntList,
    padding: impl IntList,
    dilation: impl IntList,
    groups: i64
) -> Tensor

Вот финальная программа:

Листинг 4. Прием изображения и применение операций свертки для выделения границ

use tch::{Tensor, vision::image, Kind, Device};

fn rgb_to_grayscale(tensor: &Tensor) -> Tensor {
    let red_channel = tensor.get(0);
    let green_channel = tensor.get(1);
    let blue_channel = tensor.get(2);
    
    // По формуле светового излучения вычисляется тензор оттенков серого
    let grayscale = (red_channel * 0.2989) + (green_channel * 0.5870) + (blue_channel * 0.1140);
    grayscale.unsqueeze(0)
}

fn main() {
    let mut img = image::load("mypic.jpg").expect("Failed to open image"); 
    img = rgb_to_grayscale(&img).reshape(&[1,1,1024,1024]);
    let bias: Tensor = Tensor::full(&[1], 0.0, (Kind::Float, Device::Cpu));
    
    // Определяется и применяется гауссово ядро
    let mut k1 = [-1.0, 0.0, 1.0, -2.0, 0.0, 2.0, -1.0, 0.0, 1.0];
    for element in k1.iter_mut() {
        *element /= 16.0;
    }
    let kernel1 = Tensor::from_slice(&k1)
                        .reshape(&[1,1,3,3])
                        .to_kind(Kind::Float);
    img = img.conv2d(&kernel1, Some(&bias), &[1], &[0], &[1], 1);

    // Определяется и применяется лапласово ядро
    let k2 = [0.0, 1.0, 0.0, 1.0, -4.0, 1.0, 0.0, 1.0, 0.0];
    let kernel2 = Tensor::from_slice(&k2)
                        .reshape(&[1,1,3,3])
                        .to_kind(Kind::Float);
    img = img.conv2d(&kernel2, Some(&bias), &[1], &[0], &[1], 1);


    image::save(&img, "filtered.jpg");
    
}

Вот результат операции:

В этой программе указанные ядра применены для свертки: исходное изображение преобразовано в нужные признаки, то есть границы. Далее рассмотрим реализацию этой идеи в CNN. Основное ее отличие  —  значение матрицы с ядрами выбирается сетью при обучении, то есть выбираемые из изображения признаки определяются в самой сети настройкой ядра.

Сверточный слой

Типичная структура CNN  —  это сверточные слои, за каждым из которых имеется слой подвыборки, подводимый затем обычно к полносвязным слоям. Во многом из-за слоя подвыборки, где выборка входных данных уменьшается, количество параметров сокращается.

Вот LeNet-5, одна из первых CNN:

Простая структура с полносвязным слоем уже определена в первой части. Добавим в нее определение сверточного слоя, оно пригодится при определении новых сетевых архитектур, как в листинге 1.

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

Новый сверточный слой Conv2d определяется так:

Листинг 5. Новый слой «Conv2d»

pub struct Conv2d {
    params: HashMap<String, usize>,
}

impl Conv2d {
    pub fn new (mem: &mut Memory, kernel_size: i64, in_channel: i64, out_channel: i64, stride: i64) -> Self {
        let mut p = HashMap::new();
        p.insert("kernel".to_string(), mem.new_push(&[out_channel, in_channel, kernel_size, kernel_size], true));
        p.insert("bias".to_string(), mem.push(Tensor::full(&[out_channel], 0.0, (Kind::Float, Device::Cpu)).requires_grad_(true)));
        p.insert("stride".to_string(), mem.push(Tensor::from(stride as i64)));
        Self {
            params: p,
        }
    } 
}

impl Compute for Conv2d {
    fn forward (&self,  mem: &Memory, input: &Tensor) -> Tensor {
        let kernel = mem.get(self.params.get(&"kernel".to_string()).unwrap());
        let stride: i64 = mem.get(self.params.get(&"stride".to_string()).unwrap()).int64_value(&[]);
        let bias = mem.get(self.params.get(&"bias".to_string()).unwrap());
        input.conv2d(&kernel, Some(bias), &[stride], 0, &[1], 1)
    }
}

Напомню: в структуре имеется поле params. Это коллекция типа HashMap, в типе ключа которого String хранится имя параметра, а в типе значения usize содержится местоположение конкретного параметра, то есть тензора PyTorch, в Memory, где, в свою очередь, хранятся все параметры модели.

В случае сверточного слоя в ассоциированной функции new в HashMap вставляется два параметра: Kernel и bias. При этом флагу required_gradient задается значение True. Добавляем и параметр Stride, хотя не задаем его обучаемым.

Затем реализуем в сверточном слое типаж Compute. Для этого определяется функция forward, вызываемая во время прямого прохода процесса обучения. В ней сначала с помощью метода get получаем ссылку на тензоры kernel, bias и stride из хранилища тензоров, затем вызываем функцию Conv2d  —  как в программе, только в этом случае в сети указывается используемое ядро. Жестко задан нулевой отступ, хотя этот параметр добавить не сложнее, чем Stride.

Вот и все. Это единственное добавление к начатому в первой части фреймворку для определения и обучения CNN согласно листингам 1 и 2.

Оптимизация Adam

В прошлой статье запрограммировано два алгоритма обучения: стохастического градиентного спуска и стохастического градиентного спуска с импульсом. Не менее популярен алгоритм Adam  —  закодируем на Rust и его.

В алгоритме Adam, впервые опубликованном в 2015 году, сочетаются идеи алгоритмов обучения Momentum и RMSProp:

В первой части мы реализовали тензор в Memory, а также позаботились о методах apply_grads_sgd и apply_grads_sgd_momentum  —  эквиваленте имеющейся на PyTorch функции step для градиентов. Добавляем метод new в реализацию структуры Memory, где с помощью Adam выполняется обновление градиентов:

Листинг 6. Реализация «Adam»

fn apply_grads_adam(&mut self, learning_rate: f32) {
        let mut g = Tensor::new();
        const BETA:f32 = 0.9;

        let mut velocity = Tensor::zeros(&[self.size as i64], (Kind::Float, Device::Cpu)).split(1, 0);
        let mut mom = Tensor::zeros(&[self.size as i64], (Kind::Float, Device::Cpu)).split(1, 0);
        let mut vel_corr = Tensor::zeros(&[self.size as i64], (Kind::Float, Device::Cpu)).split(1, 0);
        let mut mom_corr = Tensor::zeros(&[self.size as i64], (Kind::Float, Device::Cpu)).split(1, 0);
        let mut counter = 0;

        self.values
        .iter_mut()
        .for_each(|t| {
            if t.requires_grad() {
                g = t.grad();
                mom[counter] = BETA * &mom[counter] + (1.0 - BETA) * &g;
                velocity[counter] = BETA * &velocity[counter] + (1.0 - BETA) * (&g.pow(&Tensor::from(2)));    
                mom_corr[counter] = &mom[counter]  / (Tensor::from(1.0 - BETA).pow(&Tensor::from(2)));
                vel_corr[counter] = &velocity[counter] / (Tensor::from(1.0 - BETA).pow(&Tensor::from(2)));

                t.set_data(&(t.data() - learning_rate * (&mom_corr[counter] / (&velocity[counter].sqrt() + 0.0000001))));
                t.zero_grad();
            }
            counter += 1;
        });

}

Результаты и мнения

Как и в первой части, объективность сравнения рассмотренного кода с его эквивалентом на Python и PyTorch обеспечивается главным образом применением одинаковых гиперпараметров нейросети, алгоритмов и параметров обучения. Для тестов, запускаемых на том же ноутбуке Surface Pro 8, i7 с 16 Гб оперативной памяти и без графического процессора, применялся тот же набор данных Mnist.

В ходе многократных прогонов на Rust отмечался в среднем 60%-ный рост скорости обучения по сравнению с эквивалентом на Python. Это немало, но все же скромнее достигнутого в полносвязных сетях первой части. Меньший прирост объясняется тем, что самое затратное вычисление в CNN  —  это свертка и использовалась функция conv2d из LibTorch на C++, в целом та же, что вызывается эквивалентом на Python. Тем не менее сокращение времени на обучение модели более чем наполовину не следует игнорировать  —  это экономия многих часов, а то и дней.

Использованные материалы

Ян Гудфеллоу, Йошуа Бенджио и Аарон Курвилль, «Глубокое обучение», MIT Press, 2016 г.

Павел Карас и Давид Свобода, «Алгоритмы для эффективного вычисления свертки», в разделе «Дизайн и архитектуры цифровой обработки сигналов», Нью-Йорк, штат Нью-Йорк, США: IntechOpen, янв. 2013 г.

Янн Лекун, Леон Ботту, Йошуа Бенджио и Патрик Хаффнер, «Обучение на основе градиента, применяемое для распознавания документов», материалы Конференции IEEE 86, 2278–2324, 1998 г.

Дьедерик П. Кингма и Джимми Ба, «Adam: метод стохастической оптимизации», в материалах 3-й Международной конференции по обучению представлениям (ICLR), 2015 г.

Читайте также:

• Фича-флаги времени компиляции в Rust: зачем, как и когда используются
• Как создать API-шлюз в Rust посредством библиотеки Hyper
• Разработка макроса Rust для автоматического написания шаблонного кода SQL

Читайте нас в Telegram, VK и Дзен

---

Перевод статьи Vince Vella: Boosting Machine Learning Performance With Rust (Part 2)