Повышение производительности машинного обучения с помощью Rust (часть 2)

Эксперименты со сверточными нейронными сетями (CNN) с нуля в Rust.

В своей предыдущей статье (Часть 1) я начал свой эксперимент по разработке фреймворка машинного обучения на Rust с нуля. Основная цель моего эксперимента состояла в том, чтобы оценить улучшения скорости обучения модели, которых можно достичь, используя Rust в сочетании с PyTorch по сравнению с эквивалентом Python. Результаты оказались очень обнадеживающими для Feedforward Networks. В этой статье я продолжаю опираться на это с основной целью — иметь возможность определять и обучать сверточные нейронные сети (CNN). Как и в предыдущей статье, я продолжаю использовать крейт Tch-rs Rust в качестве оболочки библиотеки PyTorch C++ LibTorch, в основном для доступа к тензорной линейной алгебре и функциям автоградации, а остальное разрабатывается с нуля. Код частей 1 и 2 теперь доступен на 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));
}

Если попытаться сохранить определение модели как можно более похожим на эквивалент Python, приведенный выше листинг 1 должен быть интуитивно понятным для пользователей Python-PyTorch. В структуре MyModel мы теперь можем добавлять слои Conv2D, а затем инициировать их в новой связанной функции. В реализации трейта Compute определена прямая функция, которая принимает входные данные через все уровни, включая промежуточную функцию MaxPooling. В основной функции (листинг 2), как и в нашей предыдущей статье, мы обучаем нашу модель и применяем ее к набору данных Mnist.

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

Эксперименты с ядрами

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

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

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

Поиск эффективных способов вычисления сверток — очень хорошо изученная проблема (см. Ссылка). После изучения различных вариантов, которые включали несколько чистых версий ржавчины, но требовали промежуточных преобразований данных из тензоров PyTorch, я решил использовать функцию свертки LibTorch C++. Чтобы поэкспериментировать с этой функцией, я хотел создать небольшую игрушечную программу, которая берет цветное изображение, преобразует его в оттенки серого, а затем применяет некоторые известные ядра для обнаружения границ.

Сначала я попросил чат Microsoft Bing сгенерировать для меня изображение. Как только я был доволен изображением, я захотел применить функцию свертки, сначала используя ядро Гаусса, а затем ядро Лапласа.

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

Листинг 3 — метод LibTorch conv2d с открытым доступом 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);
    
    // Calculate the grayscale tensor using the luminance formula
    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));
    
    // Define and apply Gaussian Kernel
    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);

    // Define and apply Laplacian Kernel
    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 включает в себя ряд слоев свертки, за каждым из которых следует слой субдискретизации (объединение), которые затем обычно передаются в полностью связанные слои. Объединение в пул в значительной степени способствует сокращению параметров, поскольку оно снижает выборку входных данных. На приведенной ниже диаграмме изображена одна из самых ранних CNN под названием LeNet-5.

В Части 1 мы уже определили простую структуру, включающую полносвязный слой. Точно так же сейчас нам нужно внедрить определение сверточного слоя в нашу структуру, чтобы он был доступен при определении новых сетевых архитектур (как в листинге 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)
    }
}

Если вы помните мой подход из Части 1, структура содержит поле с именем params. Поле params представляет собой набор типа HashMap, где ключ имеет тип String, в котором хранится имя параметра, а значение имеет тип usize, который содержит расположение определенного параметра (который является тензором PyTorch) в нашей памяти. , который, в свою очередь, служит нашим хранилищем для всех параметров нашей модели. В случае сверточного слоя в нашей связанной функции new мы вставляем в нашу HashMap два параметра "Kernel" и "bias" с флагом required_gradient True. Я также ввожу параметр "Шаг", однако он не установлен в качестве обучаемого параметра.

Затем мы реализуем трейт Compute для нашего Convolutional Layer. Для этого необходимо определить функцию forward, которая вызывается во время прямого прохода процесса обучения. В этой функции мы сначала получаем ссылку на тензоры ядра, смещения и шага из нашего хранилища тензоров, используя метод get, а затем вызываем нашу функцию Conv2d (как мы делали в нашей игрушечной программе, однако в этом случае сеть сообщает нам, что ядро для использования). Отступы были жестко запрограммированы на нулевое значение, однако при желании его также можно легко добавить в качестве параметра, аналогичного Stride.

Вот и все! Это единственное дополнение, которое требуется в нашей небольшой структуре из Части 1, чтобы иметь возможность определять и обучать CNN, как в листинге 1–2.

Адам Оптимизация

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

Алгоритм Адама был впервые опубликован в 2015 году (Ссылка) и в основном объединяет идею алгоритмов обучения Momentum и RMSprop. Алгоритм из оригинальной статьи следующий:

В Части 1 мы реализовали нашу тензорную память, которая также обслуживает эквивалентную ступенчатую функцию градиента PyTorch (методы apply_grads_sgd и apply_grads_sgd_momentum). Следовательно, в реализацию структуры памяти добавлен новый метод, который выполняет обновление градиента с помощью Адама:

Листинг 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;
        });

}

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

Подобно моему подходу в Части 1, чтобы сравнить приведенный выше код с эквивалентом Python-PyTorch, я старался быть максимально точным, чтобы получить справедливое сравнение, в основном следя за тем, чтобы я применял одни и те же гиперпараметры нейронной сети, обучение параметры и алгоритмы обучения. Для своих тестов я также применил тот же набор данных Mnist. Я провел тесты на том же ноутбуке, Surface Pro 8, i7, с 16 ГБ ОЗУ, следовательно, без графического процессора.

После многократного запуска тестов в среднем обучение Rust привело к увеличению скорости на 60% по сравнению с эквивалентом Python. Несмотря на значительное улучшение, оно было меньше, чем в случае с FFN (мои выводы из Части 1). Я объяснил это более низкое улучшение тем фактом, что самым затратным вычислением в CNN является свертка, и, как обсуждалось выше, я решил использовать функцию LibTorch C++ conv2d, которая, в конце концов, является той же функцией, которая вызывается эквивалентом Python. Тем не менее, с учетом сказанного, сокращение времени обучения модели более чем наполовину все же не следует игнорировать — это по-прежнему означает экономию часов, если не дней!

Надеюсь, вам понравилась моя статья!

Повышение производительности машинного обучения с помощью Rust (часть 2)

Эксперименты с ядрами

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

Адам Оптимизация

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

Рекомендации

Вопросы по теме