Практическая работа №4. Автоматическая оптимизация нейронной сети с помощью Apache TVM¶

1. Цели и задачи работы¶

Цель работы – изучить программный интерфейс для автоматической оптимизации нейронных сетей с помощью Apache TVM на процессорах архитектуры RISC-V.

Достижение указанной цели предполагает решение следующих задач:

  1. Обучение архитектур логистической регрессии и полносвязной нейронной сети на наборе данных MNIST на x86-устройстве. Сохранение модели в формате Apache TVM, а также сохранение метрик качества и набора данных в формате NumPy для дальнейшего тестирования.
  2. Установка LLVM и сборка Apache TVM с LLVM.
  3. Оптимизация модели логистической регрессии.
    1. Загрузка модели логистической регрессии. Запуск, проверка корректности и измерение времени инференса без оптимизации.
    2. Оптимизация модели логистической регрессии с помощью AutoTVM, Auto-scheduler, MetaScheduler.
    3. Анализ результатов оптимизации логистической регрессии.
  4. Оптимизация полносвязной нейронной сети.
    1. Загрузка модели полносвязной нейронной сети. Запуск, проверка корректности и измерение времени инференса без оптимизации.
    2. Оптимизация модели полносвязной нейронной сети с помощью AutoTVM, Auto-scheduler, MetaScheduler.
    3. Анализ результатов оптимизации полносвязной нейронной сети.
  5. Оптимизация сверточной нейронной сети.
    1. Загрузка модели сверточной нейронной сети. Запуск, проверка корректности и измерение времени инференса без оптимизации.
    2. Оптимизация модели сверточной нейронной сети с помощью AutoTVM, Auto-scheduler, MetaScheduler.
    3. Анализ результатов оптимизации сверточной нейронной сети.

Полезные ссылки:

  • AutoTVM туториалы,
  • AutoTVM документация,
  • Auto-scheduler туториалы,
  • Auto-scheduler документация,
  • MetaScheduler документация.

Примечание: в настоящее время Apache TVM не полностью портирован на архитектуру RISC-V, в связи с этим имеется ряд ограничений, не позволяющих в полном объеме продемонстрировать имеющийся функционал для автоматической оптимизации сетей на RISC-V-устройствах. Ниже приведен примерный перечень проблем, с которыми авторы столкнулись в процессе подготовки материалов настоящей практической работы.

  • Во время оптимизации сверточных нейронных сетей возникают критические ошибки, которые не позволяют выполнить оптимизацию этих архитектур нейронных сетей на устройствах с архитектурой RISC-V. Поэтому в данной практической работе рассматриваются только полносвязные нейронные сети. В данной работе не рассматривается возможность использования оптимизации через RPC.
  • При запуске на устройствах с архитектурой RISC-V используется opt_level=2. Более высокий уровень оптимизации вызывает ошибки компиляции модели.
  • На текущий момент реализация Auto-scheduler работает с серьезными ограничениями, поэтому в данной работе Auto-scheduler используется только для оптимизации логистической регрессии.
  • Использование MetaScheduler для оптимизации приводит к критическим ошибкам, связанным с графом вычислений и ошибками компиляции.

2. Обучение моделей глубокого обучения¶

Обучение моделей выполняется на архитектуре x86 с использованием библиотеки PyTorch, так как на момент подготовки материалов работы отсутствует официальная сборка PyTorch для RISC-V-устройств.

2.1 Установка зависимостей для обучения моделей¶

2.1.2 Установка Apache TVM¶

Вначале необходимо установить Apache TVM той же версии, что будет использоваться на устройстве RISC-V, для совместимости формата хранения графа вычислений. Для этого необходимо установить LLVM и собрать Apache TVM из исходных кодов по аналогии с тем, как это было сделано в предыдущей практической работе. Ниже приведена соответствующая последовательность команд.

sudo apt install clang-17 llvm-17*
git clone --recursive https://github.com/apache/tvm
cd tvm

mkdir build
cd build

cmake -DUSE_LLVM=ON ..
make

2.1.2 Настройка окружения Python¶

Далее будем считать, что на x86-узле установлена Miniconda. Соответственно создадим виртуальное окружение для подготовки тестовых моделей. Для обучения моделей используется библиотека PyTorch и набор данных MNIST. Поэтому потребуется пакет torch, обеспечивающий функционал, необходимый для обучения/тестирования нейронных сетей, и torchvision, содержащий вспомогательные функции, в частности, для загрузки широко известных наборов данных. Далее приведена примерная последовательность команд для создания и настройки окружения.

conda create -n torch_train python==3.10
conda activate torch_train

pip install numpy matplotlib torchmetricspip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install notebook

2.2 Обучение моделей¶

Перед обучением модели необходимо активировать созданное на предыдущем этапе виртуальное окружение и установить путь к Apache TVM.

conda activate torch_train
export PYTHONPATH=<PATH TO TVM>/python:${PYTHONPATH}

Процесс обучения реализован в файле 04_train_model_x86.ipynb. Более подробно возможности библиотеки PyTorch для обучения моделей рассматривались во второй практической работе. Необходимо запустить выполнение этого файла. После завершения его работы архитектура и веса обученных нейронных сетей будут сохранены в файл в директории model/. В этой же директории будет сохранен файл с показателями точности моделей. Наряду с этим, указанный скрипт обеспечивает сохранение тестовых данных (изображения и их метки) для упрощения процедуры их загрузки на RISC-V-устройствах. Соответственное данные сохраняются в директорию data/.

3. Сборка и установка LLVM и Apache TVM¶

3.1. Сборка LLVM¶

В данной практической работе не используется кросс-компиляция моделей или слоев. Компиляция происходит на устройстве с архитектурой RISC-V. Поэтому требуется собрать Apache TVM с LLVM. Рекомендуется использовать версию 15 <= LLVM <= 17.

3.1.1. Установка с помощью менеджера пакетов¶

sudo apt install clang-17 llvm-17*

3.1.1. Сборка LLVM версии llvmorg-17.0.6 (для версии llvmorg-17.0.6)¶

Для сборка LLVM из исходных кодов требуется загрузить необходимую версию LLVM из репозитори GitHub. В данном работе используется версия llvmorg-17.0.6, далее, используя утилиту CMake сгенерировать make-файлы и выполнить сборку. Ниже приведена соответствующая последовательност команд.

git clone https://github.com/llvm/llvm-project.git -b llvmorg-17.0.6
cd llvm-project

mkdir _build
cd _build

cmake -DCMAKE_BUILD_TYPE="Release" \
      -DLLVM_ENABLE_PROJECTS=clang \
      -DBUILD_SHARED_LIBS=True \
      -DLLVM_USE_SPLIT_DWARF=True \
      -DCMAKE_INSTALL_PREFIX="../../_install" ../llvm

make

Примечание: в случае сборки LLVM из исходных кодов перед сборкой Apache TVM необходимо указать путь к LLVM в переменной окружения PATH и создать переменную окружения LLVM_CONFIG. Ниже показан пример.

PATH="<PATH TO LLVM>/_build/bin:$PATH"
export LLVM_CONFIG=<PATH TO LLVM>/_build/bin/llvm-config

3.2. Установка OpenBLAS¶

Далее необходимо установить OpenBLAS, используя менеджер пакетов.

sudo apt-get install libopenblas-dev

3.3. Настройка окружения Python¶

Для выполнения практической работы создадим и настроим виртуальное окружение Python так, как показано ниже:

python3 -m venv ~/tvm_cpu/
source ~/tvm_cpu/bin/activate

pip install scipy numpy matplotlib pandas
pip install cloudpickle traitlets typing-extensions psutil pybind11 decorator attrs 
pip install notebook

3.4. Сборка Apache TVM¶

Для сборки Apache TVM используем ветку main GitHub-репозитория, так как недавно были внесены критически важные исправленияя 1 и 2. Для сборки Apache TVM не обязательн использовать созданную виртуальную среду для Python.

git clone --recursive https://github.com/apache/tvm
cd tvm

mkdir build
cd build

cmake -DUSE_LLVM=ON -DUSE_BLAS=openblas ..
make

3.5. Активация окружения для практической работы¶

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

source ~/tvm_cpu/bin/activate
export PYTHONPATH=<PATH TO TVM>/python:${PYTHONPATH}

4. Программная реализация вспомогательных функций¶

4.1. Импорт пакетов¶

Для использования функционала Apache TVM и других вспомогательных библиотек импортируем необходимые пакеты.

Также определим переменную, содержащую используемый тип данных для элементов тензоров - float32, а также установим в качестве целевого устройства для запуска CPU.

In [1]:
import os
from time import time

import matplotlib.pyplot as plt


import numpy as np
import tvm

from tvm import autotvm
from tvm import auto_scheduler
from tvm import meta_schedule as ms

from tvm import relay
from tvm.autotvm.tuner import XGBTuner
from tvm.contrib import graph_executor


dtype = 'float32'
dev = tvm.cpu()

global_trial = 96

4.2. Строка компиляции¶

На данном этапе определим строку компиляции target. Компиляция нейронных сетей происходит на устройстве с архитектурой RISC-V без использования кросс-компиляции. Для упрощения тестирования и отладки добавлена возможность запуска на x86_64.

Для определения архитектуры устройства необходимо создать обьект строки компиляции по умолчанию для LLVM - tvm.target.Target('llvm'). Далее с помощью атрибута mtriple выбрать строку компиляции:

  • Если атрибут mtriple отсутствует или содержит подстроку x86_64, используется стандартная строка компиляции llvm.
  • Если mtriple содержит подстроку riscv64, используется строка компиляции llvm -jit=orcjit -mtriple=riscv64-unknown-linux-gnu -mcpu=generic-rv64 -mabi=lp64d -mattr=+64bit,+m,+a,+f,+d.
  • В противном случае генерируется исключение.

Примечание 1: если TVM устанавливался через PyPI, то mtriple пустой.

Примечание 2: TVM поддерживает различные бэкенды, такие, как llvm, opencl, cuda и прочие. В данном случае для генерации машинного кода TIR будет транслироваться в LLVM IR, после чего из LLVM IR будет генерироваться машинный код. Краткое описание параметров строки компиляции приведено ниже.

  • -jit=orcjit указывает на использование JIT-компилятора ORC (On-Request Compilation). TVM необходим данный ключ при компиляции на RISC-V.

  • -mtriple=riscv64-unknown-linux-gnu определяет тройку целевой архитектуры. Она указывает на платформу RISC-V 64-бит с операционной системой Linux и неуточненным вендором.

  • -mcpu=generic-rv64 указывает целевой тип процессора.

  • -mabi=lp64d определяет используемый ABI (Application Binary Interface). lp64d обозначает ABI, в котором длинные целые (long) и указатели (pointers) имеют размер 64 бита, и включена поддержка вещественных чисел двойной точности (d).

  • -mattr=+64bit,+m,+a,+f,+d задает атрибуты целевой архитектуры.

    • +64bit - поддержка 64-битной архитектуры.
    • +m - поддержка умножения и деления.
    • +a - поддержка атомарных операций.
    • +f - поддержка операций с плавающей запятой одинарной точности.
    • +d - поддержка операций с плавающей запятой двойной точности.

4.3. Уровень оптимизации графа¶

При запусках на устройствах RISC-V используется opt_level=2, в случае запуска на архитектуре x86-64 используется opt_level=3.

In [2]:
def is_x86():
    if tvm.target.Target('llvm').attrs.get('mtriple') is None:
        return True
    return 'x86_64' in tvm.target.Target('llvm').attrs.get('mtriple')

def is_riscv():
    return 'riscv64' in tvm.target.Target('llvm').attrs.get('mtriple')
    

print(f"mtriple устройства {tvm.target.Target('llvm').attrs.get('mtriple')}")

if is_x86():
    target = tvm.target.Target('llvm')
    opt_level = 3
elif is_riscv():
    target = tvm.target.Target(
        'llvm -jit=orcjit -mtriple=riscv64-unknown-linux-gnu '
        '-mcpu=generic-rv64 -mabi=lp64d -mattr=+64bit,+m,+a,+f,+d'
    )
    opt_level = 2
else:
    raise ValueError("Unsupported architecture")

print(f'{target = }')
mtriple устройства riscv64-unknown-linux-gnu
target = llvm -keys=cpu -jit=orcjit -mabi=lp64d -mattr=+64bit,+m,+a,+f,+d -mcpu=generic-rv64 -mtriple=riscv64-unknown-linux-gnu

4.3 Вспомогательные функции¶

Реализуем функцию load_model для загрузки модели в формате TVM, а также функцию load_images_and_labels для загрузки изображений и меток из набора данных MNIST.

In [3]:
def load_model(mod_file, params_file):
    with open(mod_file, "r") as fo:
        mod = fo.read()
        
    mod = tvm.ir.load_json(mod)

    with open(params_file, "rb") as fo:
        params = relay.load_param_dict(fo.read())
    
    return mod, params

def load_images_and_labels(images_path, labels_path):
    images = np.load(images_path)
    labels = np.load(labels_path)
    
    return images, labels

Далее выполним реализацию функции timeit_inference для измерения времени инференса и функции get_accuracy для определения качества решения задачи.

  1. Функция timeit_inference. Измерение времени инференса проводится на наборе данных MNIST. Инференс выполняется отдельно для каждого изображения из набора данных MNIST. Время выполнения и результаты предсказания (номер класса, на котором достигается максимумальная достоверность) возвращаются из функции.
  2. Функция get_accuracy. Определение качества решения задачи выполняется для всего набора данных MNIST посредством сравнения результатов предсказания и разметки. Точность вычисляется как отношение количества совпадений предсказанных и размеченных классов к общему числу изображений в наборе данных.
In [4]:
def timeit_inference(mod, lib, images):
    input_name = mod['main'].params[0].name_hint
    input_shape = mod['main'].params[0].type_annotation.shape
    input_shape = [int(s) for s in input_shape]

    dev = tvm.cpu()
    module = graph_executor.GraphModule(lib["default"](dev))
    
    predict = []
    times = []
    for i in range(len(images)):
        img = np.array(images[i:i+1], dtype=np.float32).reshape(input_shape)
        module.set_input(input_name, img)

        ts = time()
        module.run()
        tf = time()
        times.append((tf - ts) * 1000)
        
        output = module.get_output(0).numpy()
        predict.append(np.argmax(output))
        
    return np.array(predict), np.array(times)
    
def get_accuracy(labels, predict):
    return np.mean(labels == predict)

На данном этапе необходимо загрузить изображения, разметку и информацию о точности работы нейронных сетей, полученную после обучения на системе с архитектурой x86-64. Далее при решении задач практической работы точность нейронной сети необходимо сопоставлять с загруженными значениями.

In [5]:
images, labels = load_images_and_labels('data/test_images.npy', 'data/test_labels.npy')

metric = np.load('model/metric.npy', allow_pickle='TRUE').item()
print(metric)
{'logreg': array(0.9264, dtype=float32), 'fcnn': array(0.9804, dtype=float32), 'cnn': array(0.985, dtype=float32)}

5. Общая информация про методы автоматической оптимизации слоев в Apache TVM¶

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

Ключевым параметром в процессе оптимизации является количество итераций оптимизации задач. Подбор этого параметра является нетривиальной задачей:

  • Если значение параметра слишком маленькое, эффективное решение может не быть найдено.
  • Слишком большое значение параметра приведет к значительным затратам времени.
  • Оптимальное количество итераций зависит от характеристик нейронной сети, целевого устройства и используемого метода оптимизации.

На каждой итерации выполняется несколько проверок качества конкретной реализации. Параметры этих проверок задаются через обьекты классов autotvm.measure_option, auto_scheduler.LocalRunner и ms.runner.LocalRunner, которые предоставляют интерфейс для указания числа замеров производительности:

  • number - количество запусков кода для усреднения времени выполнения в процессе одного замера.
  • repeat - число замеров. Всего выполняется (1 + number x repeat) запусков, где первый запуск используется для прогрева и не учитывается.
  • enable_cpu_cache_flush очищает кэш CPU между последовательными замерами для более точной оценки задержек.

Таким образом, чем больше значение number x repeat, тем более точной будет оценка времени работы планов вычислений, однако, это также увеличивает продолжительность процесса автоматической оптимизации.

Примечание: псевдокод работы методов оценки времени выполнения в Apache TVM приведен ниже.

for r in range(repeat):
    time_start = now()
    for n in range(number):
        func_name()
    time_end = now()
    total_times.append((time_end - time_start) / number)

6. Запуск и оптимизация модели логистической регрессии¶

6.1. Компиляция и запуск модели¶

Вначале необходимо выполнить загрузку модели логистической регрессии.

In [6]:
default_logreg_time, autotvm_logreg_time, autoscheduler_logreg_time, ms_logreg_time = 0, 0, 0, 0

mod, params = load_model('model/logreg.json', 'model/logreg.params')
print(mod['main'])
fn (%input0: Tensor[(1, 784), float32] /* span=aten::linear_0.input0:0:0 */, %aten::linear_0.weight: Tensor[(10, 784), float32] /* span=aten::linear_0.weight:0:0 */, %aten::linear_0.bias: Tensor[(10), float32] /* span=aten::linear_0.bias:0:0 */) {
  %0 = nn.dense(%input0, %aten::linear_0.weight, units=None) /* span=aten::linear_0:0:0 */;
  nn.bias_add(%0, %aten::linear_0.bias, axis=-1) /* span=aten::linear_0:0:0 */
}

Следующий шаг - компиляция модели без оптимизации слоев.

In [7]:
with tvm.transform.PassContext(opt_level=opt_level):
    lib = relay.build(mod, target=target, params=params)
One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.

После компиляции можно выполнить запуск вывода и измерение времени выполнения с использованием разработанной функции timeit_inference, а также определение качества работы логистической регрессии с помощью функции get_accuracy и сравнение полученной точности классификации с загруженным значением, которое получено на x86-64.

In [8]:
default_logreg_predict, default_logreg_times = timeit_inference(mod, lib, images)

default_logreg_accuracy = get_accuracy(labels, default_logreg_predict)
assert np.allclose(metric['logreg'], default_logreg_accuracy, rtol=1e-5)

default_logreg_time = np.median(default_logreg_times)
print(f'Медианное время работы неоптимизированной модели: {default_logreg_time:.4f} мc')
Медианное время работы неоптимизированной модели: 0.1159 мc

6.2. Использованием возможностей AutoTVM¶

Определим функцию get_autotvm_task для извлечения задач и вывода информации о задачах (номер задачи и task.workload). Для этого используем метод autotvm.task.extract_from_program, передав на вход модель, целевое устройство и обученные параметры модели. В данном случае рассматриваются два типа задач: полносвязные слои без трансформации весов и с трансформацией весов для улучшения работы с памятью.

Для архитектур x86 и RISC-V задачи обозначаются как dense_*.x86. На данный момент в Apache TVM нет реализаций планов вычислений для RISC-V. Благодаря тому, что Apache TVM опирается на возможности LLVM в процессе компиляции и сходство архитектур, инструмент успешно использует планы вычислений, разработанные для x86-платформ, на устройствах с архитектурой RISC-V.

In [9]:
def get_autotvm_task(
    mod: tvm.ir.module.IRModule, 
    target: tvm.target.target.Target, 
    params: tvm.ir.container.Map
) -> list[tvm.autotvm.task.task.Task, ...]:
    """
    Параметры:
        mod: Модуль IRModule.
        target: Строка компиляции.
        params: Веса нейронной сети.
    
    Возвращаемое значение:
        Список задач.
    """
    print("Извлечение задач\n")

    tasks = autotvm.task.extract_from_program(
        mod, target=target, params=params,
    )

    for idx, task in enumerate(tasks):
        print(f"Номер задачи: {idx}\nИнформация о задаче: {task.workload}\n")
        
    return tasks

Вызовем разработанную функцию get_autotvm_task для извлечения задач из графа вычислений для AutoTVM.

In [10]:
tasks = get_autotvm_task(mod, target, params)
Извлечение задач

Номер задачи: 0
Информация о задаче: ('dense_nopack.x86', ('TENSOR', (1, 784), 'float32'), ('TENSOR', (10, 784), 'float32'), None, 'float32')

Номер задачи: 1
Информация о задаче: ('dense_pack.x86', ('TENSOR', (1, 784), 'float32'), ('TENSOR', (10, 784), 'float32'), None, 'float32')

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

Вначале необходимо определить параметры проверки времени выполнения каждого плана с помощью autotvm.measure_option и autotvm.LocalRunner. Затем для каждой задачи определить модель затрат. В качестве модели затрат для оценки времени выполнения слоя используется метод градиентного бустинга деревьев, реализованный на базе XGBoost. Apache TVM предоставляет интерфейс для нескольких методов оптимизации. Инициализируем для каждой задачи класс XGBTuner.

Каждая задача оптимизируется min(n_trial, len(task.config_space)) раз, где n_trial - заданное количество попыток, а len(task.config_space) - количество различных конфигураций в плане вычислений для данного тензорного выражения.

После определения всех параметров необходимо запустить оптимизацию с помощью метода tuner_obj.tune, передав в качестве параметров количество экспериментов оптимизации для каждой задачи, объект measure_option и название файла для логирования через autotvm.callback.log_to_file(log_file).

In [11]:
def tune_autotvm(
    tasks: list[tvm.autotvm.task.task.Task, ...], 
    n_trial: int, 
    log_file: str
):
    """
    Параметры:
        tasks: Список задач.
        n_trial: Количество экспериментов для каждой задачи.
        log_file: Файл для логирование результатов оптимизации.
    """    
    measure_option = autotvm.measure_option(
        builder=autotvm.LocalBuilder(),
        runner=autotvm.LocalRunner(repeat=1, number=3, enable_cpu_cache_flush=True),
    )

    for i, task in enumerate(tasks):
        prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))
        tuner_obj = XGBTuner(task)

        n = min(n_trial, len(task.config_space))

        tuner_obj.tune(
            n_trial=n,
            measure_option=measure_option,
            callbacks=[
                autotvm.callback.progress_bar(n_trial, prefix=prefix),
                autotvm.callback.log_to_file(log_file),
            ],
        )

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

In [12]:
os.makedirs('autotvm/', exist_ok=True)
log_file = 'autotvm/autotvm_logreg.log'
n_trial = global_trial

tune_autotvm(tasks, n_trial, log_file)
[Task  2/ 2]  Current/Best:    0.14/   1.68 GFLOPS | Progress: (60/96) | 103.98 s Done.
[Task  2/ 2]  Current/Best:    0.62/   1.68 GFLOPS | Progress: (96/96) | 152.62 s Done.

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

In [13]:
with autotvm.apply_history_best(log_file):
    with tvm.transform.PassContext(opt_level=opt_level):
        lib = relay.build(mod, target=target, params=params)

На данном этапе можно выполнить измерение времени вывода с использованием функции timeit_inference, проверку качества работы оптимизированной модели с помощью функции get_accuracy и сравнение точности классификации с рефенсным значением, которое было получено после запуска обучения модели.

In [14]:
autotvm_logreg_predict, autotvm_logreg_times = timeit_inference(mod, lib, images)

autotvm_logreg_accuracy = get_accuracy(labels, autotvm_logreg_predict)
assert np.allclose(metric['logreg'], autotvm_logreg_accuracy, rtol=1e-5)

autotvm_logreg_time = np.median(autotvm_logreg_times)
print(f'Медианное время работы после оптимизации слоев с помощью AutoTVM: {autotvm_logreg_time:.4f} мc')
Медианное время работы после оптимизации слоев с помощью AutoTVM: 0.0393 мc

6.3. Использование Auto-scheduler¶

Определим функцию get_auto_scheduler_task для извлечения задач и вывода информации о задачах (номер задачи и task.desc).

Аналогично AutoTVM, вначале необходимо извлечь задачи, используя метод auto_scheduler.extract_tasks, передав в качестве входных параметров модель, целевое устройство для запуска вывода, набор обученных параметров модели. Также Auto-scheduler позволяет регулировать уровень оптимизации графа с помощью параметра opt_level. Это значение должно совпадать с уровнем оптимизации графа вычислений при компиляции модели. Отметим, что в данном случае, граф вычислений состоит только из одного слоя, поэтому объединение слоев не будет выполняться. Метод возвращает значение task_weights, которое определяет вес каждого подграфа. По умолчанию вес равен $1$. Если присутствуют $N$ одинаковых подграфов, то они будут представлены в виде одной задачи с весом $N$.

In [15]:
def get_auto_scheduler_task(
    mod: tvm.ir.module.IRModule, 
    target: tvm.target.target.Target, 
    params: tvm.ir.container.Map,
    opt_level: int
) -> tuple[list[tvm.auto_scheduler.search_task.SearchTask, ...], list[int, ...]]:
    """
    Параметры:
        mod: Модуль IRModule.
        target: Строка компиляции.
        params: Веса нейронной сети.
        opt_level: Уровень оптимизации графа вычислений.
    
    Возвращаемое значение:
        Список задач и список весов задач.
    """
    tasks, task_weights = auto_scheduler.extract_tasks(mod, target=target, params=params, opt_level=opt_level)

    for idx, task in enumerate(tasks):
        print(f"Номер задачи: {idx}\nИнформация о задаче: {task.desc}\n")
        
    return tasks, task_weights

Выполним извлечение задач для Auto-scheduler, вызвав функцию get_auto_scheduler_task.

In [16]:
tasks, task_weights = get_auto_scheduler_task(mod, target, params, opt_level)
Номер задачи: 0
Информация о задаче: vm_mod_fused_nn_dense_nn_bias_add

Далее реализуем функцию tune_auto_scheduler для автоматической настройки параметров оптимизации нейронной сети.

В данном случае для оптимизации необходимо создать обьект класса auto_scheduler.TaskScheduler с описанием задач и определить параметры оптимизации auto_scheduler.TuningOptions. После этого можно вызвать метод tune для созданного объекта класса auto_scheduler.TaskScheduler.

Примечания:

  1. При определении параметров оптимизации используется параметр num_measures_per_round. Он определяет количество конфигураций аннотированных эскизов, для которых будет измерено время перед обновлением базы результатов. После обновления базы результатов модель затрат переобучается, и запускается новая итерация эволюционного алгоритма для генерации новых эскизов.
  2. Параметр количества оптимизаций num_measure_trials в Auto-scheduler задает общее количество измерений для всех подграфов.
In [17]:
def tune_auto_scheduler(
    tasks: list[tvm.auto_scheduler.search_task.SearchTask, ...], 
    task_weights: list[int, ...], 
    log_file: str, 
    n_trials: int
):
    """
    Параметры:
        tasks: Список задач.
        task_weights: Список весов задач.
        n_trial: Количество экспериментов для каждой задачи.
        log_file: Файл для логирования результатов оптимизации.
    """
    tuner = auto_scheduler.TaskScheduler(tasks, task_weights, strategy='round-robin')
    tune_option = auto_scheduler.TuningOptions(
        num_measure_trials=n_trials,
        num_measures_per_round=8,
        runner=auto_scheduler.LocalRunner(repeat=1, number=3, enable_cpu_cache_flush=True),
        measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
        verbose=1,
    )

    tuner.tune(tune_option)

На данном этапе можно выполнить запуск оптимизации с помощью Auto-scheduler. Определим файл с навзанием log_file для логирования результатов оптимизации. Установим число экспериментов при оптимизации равным N * len(tasks). Выполним запуск оптимизации посредством вызова функции tune_auto_scheduler.

In [18]:
os.makedirs('auto_schedule/', exist_ok=True)
log_file = 'auto_schedule/auto-schedule_logreg.log'
n_trial_per_task = global_trial

tune_auto_scheduler(tasks, task_weights, log_file, n_trial_per_task * len(tasks))
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------

-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |            - |              - |      0 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: - ms	Trials: 0	Used time : 0 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Generate Sketches		#s: 5
Sample Initial Population	#s: 917	fail_ct: 686	Time elapsed: 11.80
GA Iter: 0	Max score: 0.9981	Min score: 0.9818	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.9998	Min score: 0.9980	#Pop: 16	#M+: 1392	#M-: 74
EvolutionarySearch		#s: 16	Time elapsed: 52.61
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
.E.E.E.....*****
Time elapsed for measurement: 19.50 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.35 s
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.091 |           0.17 |      8 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.091 ms	Trials: 8	Used time : 84 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 949	fail_ct: 691	Time elapsed: 11.92
GA Iter: 0	Max score: 0.9982	Min score: 0.9827	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.9997	Min score: 0.9979	#Pop: 16	#M+: 1383	#M-: 79
EvolutionarySearch		#s: 16	Time elapsed: 52.49
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
.E.E.E.....*****
Time elapsed for measurement: 17.59 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.35 s
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.091 |           0.17 |     16 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.091 ms	Trials: 16	Used time : 167 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 925	fail_ct: 675	Time elapsed: 11.77
GA Iter: 0	Max score: 0.9939	Min score: 0.9176	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.9943	Min score: 0.9377	#Pop: 16	#M+: 1374	#M-: 79
EvolutionarySearch		#s: 16	Time elapsed: 55.13
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 19.95 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.35 s
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.055 |           0.29 |     24 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.055 ms	Trials: 24	Used time : 254 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 939	fail_ct: 681	Time elapsed: 11.98
GA Iter: 0	Max score: 0.6244	Min score: 0.5429	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.9361	Min score: 0.8612	#Pop: 16	#M+: 1379	#M-: 68
EvolutionarySearch		#s: 16	Time elapsed: 54.20
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 21.73 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.35 s
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.055 |           0.29 |     32 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.055 ms	Trials: 32	Used time : 343 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 925	fail_ct: 680	Time elapsed: 11.89
GA Iter: 0	Max score: 0.6703	Min score: 0.5938	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.9173	Min score: 0.8804	#Pop: 16	#M+: 1369	#M-: 68
EvolutionarySearch		#s: 16	Time elapsed: 54.50
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 21.17 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.39 s
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.053 |           0.29 |     40 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.053 ms	Trials: 40	Used time : 431 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 892	fail_ct: 699	Time elapsed: 11.70
GA Iter: 0	Max score: 0.6277	Min score: 0.5668	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.9677	Min score: 0.8433	#Pop: 16	#M+: 1383	#M-: 66
EvolutionarySearch		#s: 16	Time elapsed: 54.26
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........T*******
Time elapsed for measurement: 32.01 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.67 s
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.017 |           0.92 |     48 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.017 ms	Trials: 48	Used time : 530 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 915	fail_ct: 682	Time elapsed: 11.77
GA Iter: 0	Max score: 0.4036	Min score: 0.3426	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.8270	Min score: 0.6361	#Pop: 16	#M+: 1386	#M-: 52
EvolutionarySearch		#s: 16	Time elapsed: 54.76
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 21.11 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.42 s
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.017 |           0.92 |     56 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.017 ms	Trials: 56	Used time : 618 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 919	fail_ct: 703	Time elapsed: 11.89
GA Iter: 0	Max score: 0.3851	Min score: 0.3053	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.9138	Min score: 0.7792	#Pop: 16	#M+: 1388	#M-: 42
EvolutionarySearch		#s: 16	Time elapsed: 56.43
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 23.07 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.43 s
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.016 |           1.00 |     64 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.016 ms	Trials: 64	Used time : 710 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 946	fail_ct: 688	Time elapsed: 12.03
GA Iter: 0	Max score: 0.3369	Min score: 0.2783	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.9011	Min score: 0.6845	#Pop: 16	#M+: 1386	#M-: 34
EvolutionarySearch		#s: 16	Time elapsed: 57.43
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 21.04 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.42 s
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.013 |           1.24 |     72 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.013 ms	Trials: 72	Used time : 801 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 936	fail_ct: 661	Time elapsed: 11.88
GA Iter: 0	Max score: 0.2867	Min score: 0.2397	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.8073	Min score: 0.5653	#Pop: 16	#M+: 1394	#M-: 30
EvolutionarySearch		#s: 16	Time elapsed: 57.72
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 18.06 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.40 s
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.012 |           1.31 |     80 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.012 ms	Trials: 80	Used time : 889 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 924	fail_ct: 699	Time elapsed: 11.95
GA Iter: 0	Max score: 0.3326	Min score: 0.2915	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.8679	Min score: 0.6444	#Pop: 16	#M+: 1389	#M-: 25
EvolutionarySearch		#s: 16	Time elapsed: 57.49
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 18.49 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.41 s
----------------------------------------------------------------------
------------------------------  [ Task Scheduler ]
----------------------------------------------------------------------
|  ID  |                       Task Description                        | Latency (ms) | Speed (GFLOPS) | Trials |
-----------------------------------------------------------------------------------------------------------------
|    0 |                             vm_mod_fused_nn_dense_nn_bias_add |        0.011 |           1.47 |     88 |
-----------------------------------------------------------------------------------------------------------------
Estimated total latency: 0.011 ms	Trials: 88	Used time : 978 s	Next ID: 0	
----------------------------------------------------------------------
------------------------------  [ Search ]
----------------------------------------------------------------------
Sample Initial Population	#s: 944	fail_ct: 653	Time elapsed: 11.72
GA Iter: 0	Max score: 0.3208	Min score: 0.2541	#Pop: 16	#M+: 0	#M-: 0
GA Iter: 4	Max score: 0.6711	Min score: 0.5192	#Pop: 16	#M+: 1392	#M-: 27
EvolutionarySearch		#s: 16	Time elapsed: 57.87
----------------------------------------------------------------------
------------------------------  [ Measure ]
----------------------------------------------------------------------
Get 8 programs to measure:
........********
Time elapsed for measurement: 20.90 s
----------------------------------------------------------------------
------------------------------  [ Train cost model ]
----------------------------------------------------------------------
Time elapsed for training: 0.41 s

По завершении оптимизации необходимо скомпилировать модель с учетом истории оптимизации.

In [19]:
with auto_scheduler.ApplyHistoryBest(log_file):
    with tvm.transform.PassContext(
        opt_level=opt_level, config={"relay.backend.use_auto_scheduler": True},
    ):
        lib = relay.build(mod, target=target, params=params)

Далее для скомпилированной модели можно выполнить измерение времени выполнения с использованием вызова функции timeit_inference, определить качество работы с помощью функции get_accuracy и проверить корректность, сравнив полученное значение показателя точности с референсным значением.

In [20]:
autoscheduler_logreg_predict, autoscheduler_logreg_times = timeit_inference(mod, lib, images)

autoscheduler_logreg_accuracy = get_accuracy(labels, autoscheduler_logreg_predict)
assert np.allclose(metric['logreg'], autoscheduler_logreg_accuracy, rtol=1e-5)

autoscheduler_logreg_time = np.median(autoscheduler_logreg_times)
print(f'Медианное время работы после оптимизации слоев с помощью Auto-scheduler: {autoscheduler_logreg_time:.4f} мc')
Медианное время работы после оптимизации слоев с помощью Auto-scheduler: 0.0417 мc

6.4. Применение MetaScheduler¶

Использование MetaScheduler требует указания числа ядер при формировании строки, содержащей параметры целевого устройства, например, -num-cores 4. Данный параметр можно указать равным количеству физических ядер на устройстве. Внесем соответствующие изменения в исходный код.

In [21]:
print(f"mtriple устройства {tvm.target.Target('llvm').attrs.get('mtriple')}")

if is_x86():
      target = tvm.target.Target('llvm -num-cores 6')
elif is_riscv():    
    target = tvm.target.Target(
        'llvm -jit=orcjit -mtriple=riscv64-unknown-linux-gnu '
        '-mcpu=generic-rv64 -mabi=lp64d -mattr=+64bit,+m,+a,+f,+d -num-cores 4'
    )
else:
    raise ValueError("Unsupported architecture")


    print(f'{target = }')
mtriple устройства riscv64-unknown-linux-gnu

Определим функцию get_ms_task для извлечения задач и вывода информации о задачах (номер задачи и task.task_name). Аналогично предыдущим методам оптимизации, извлечение задач выполняется с помощью методов ms.relay_integration.extract_tasks и ms.relay_integration.extracted_tasks_to_tune_contexts.

In [22]:
def get_ms_task(
    mod: tvm.ir.module.IRModule, 
    target: tvm.target.target.Target, 
    params: tvm.ir.container.Map,
    opt_level: int,
    work_dir: str
) -> tuple[list[tvm.meta_schedule.tune_context.TuneContext, ...], list[int, ...]]:
    """
    Параметры:
        mod: Модуль IRModule.
        target: Строка компиляции.
        params: Веса нейронной сети.
        opt_level: Уровень оптимизации графа вычислений.
        work_dir: Директория для логирования результатов оптимизации.
    
    Возвращаемое значение:
        Список задач и список весов задач.
    """
    extracted_tasks = ms.relay_integration.extract_tasks(
        mod, target=target, params=params, opt_level=opt_level,
    )

    tasks, task_weights = ms.relay_integration.extracted_tasks_to_tune_contexts(
        extracted_tasks, work_dir
    )

    for idx, task in enumerate(tasks):
        print(f"Номер задачи: {idx}\nИнформация о задаче: {task.task_name}\n")
        
    return tasks, task_weights

Вызовем разработанную функцию get_ms_task, предварительно определив директорию work_dir для логирования результатов оптимизации.

In [23]:
work_dir = "meta_schedule_logreg"

if is_x86():
    tasks, task_weights = get_ms_task(mod, target, params, opt_level, work_dir)

По аналогии с другими рассмотренными методами реализуем функцию tune_ms для автоматической настройки параметров запуска вывода нейронной сети. Данная функция должна вызывать метод ms.tune.tune_tasks, который принимает на вход набор задач, веса этих задач и параметры оптимизации.

Примечание: для указания количества запусков при оценке качества эскиза на вход ms.tune.tune_tasks передается объект ms.runner.LocalRunner с указанием параметра ms.runner.config.EvaluatorConfig.

In [24]:
def tune_ms(
    tasks: list[tvm.meta_schedule.tune_context.TuneContext, ...], 
    task_weights: list[int, ...], 
    work_dir: str, 
    n_trials: int
):
    """
    Параметры:
        tasks: Список задач.
        task_weights: Список весов задач.
        work_dir: Директория для логирования результатов оптимизации.
        n_trial: Количество экспериментов для каждой задачи.
    """
    
    if not os.path.exists(work_dir):
        os.mkdir(work_dir)

    ms.tune.tune_tasks(
        tasks=tasks,
        task_weights=task_weights,
        work_dir=work_dir,
        max_trials_global=n_trials,
        num_trials_per_iter=8,
        builder=ms.builder.LocalBuilder(),
        runner=ms.runner.LocalRunner(
            evaluator_config=ms.runner.config.EvaluatorConfig(repeat=1, number=3, enable_cpu_cache_flush=True)
        ),
    )

Далее выполним запуск оптимизации с помощью MetaScheduler посредством вызова функции tune_ms, установив число экспериментов при оптимизации равным N * len(tasks).

In [25]:
n_trial_per_task = global_trial

if is_x86():
    tune_ms(tasks, task_weights, work_dir, n_trial_per_task * len(tasks))

После оптимизации можно скомпилировать нейронную с учетом построенных оптимизаций с помощью интерфейса MetaScheduler ms.relay_integration.compile_relay.

In [26]:
if is_x86():
    
    database = ms.database.JSONDatabase(
        f"{work_dir}/database_workload.json",
        f"{work_dir}/database_tuning_record.json",
        allow_missing=False
    )

    lib = ms.relay_integration.compile_relay(
        database, mod, target, params,
        opt_level=opt_level,
    )

В завершении измерим время вывода с использованием функции timeit_inference, определим качество работы модели с помощью функции get_accuracy и выполним проверку корректности работы оптимизированной модели, сравнив полученное значение показателя точности с референсным.

In [27]:
if is_x86():
    
    ms_logreg_predict, ms_logreg_times = timeit_inference(mod, lib, images)

    ms_logreg_accuracy = get_accuracy(labels, ms_logreg_predict)
    assert np.allclose(metric['logreg'], ms_logreg_accuracy, rtol=1e-5)

    ms_logreg_time = np.median(ms_logreg_times)
    print(f'Медианное время работы после оптимизации слоев с помощью MetaScheduler: {ms_logreg_time:.4f} мc')

6.5. Анализ полученных результатов¶

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

In [28]:
fig, ax = plt.subplots()

name = ['Без оптимизации\nслоев', 'AutoTVM', 'Auto-scheduler', 'MetaScheduler']
times = [default_logreg_time, autotvm_logreg_time, autoscheduler_logreg_time, ms_logreg_time]
bar_labels = ['red', 'blue', '_red', 'orange']
bar_colors = ['tab:blue', 'tab:red', 'tab:green', 'tab:orange']

bars = ax.bar(name, times, label=name, color=bar_colors)
ax.set_title('Среднее время\nвыполнения (мс)', fontsize=18)

for bar, n, t in zip(bars, name, times):
    h = bar.get_height()
    if n == 'Без оптимизации\nслоев': h = h / 2
    if h != 0:
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            h,
            f'{round(t, 4)} с',
            ha='center',
            va='bottom',
            fontsize=15,
        )

ax.xaxis.label.set_size(40)
ax.set_title('Среднее время\nвыполнения (с)', fontsize=18)
plt.grid()
No description has been provided for this image

Вывод: оптимизация значительно ускоряет время работы сети.

7. Запуск и оптимизация полносвязной нейронной сети¶

7.1. Компиляция и запуск модели¶

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

In [29]:
default_fcnn_time, autotvm_fcnn_time, ms_fcnn_time = 0, 0, 0

mod, params = load_model('model/fcnn.json', 'model/fcnn.params')
print(mod['main'])
fn (%input0: Tensor[(1, 784), float32] /* span=aten::linear_0.input0:0:0 */, %aten::linear_0.weight: Tensor[(300, 784), float32] /* span=aten::linear_0.weight:0:0 */, %aten::linear_0.bias: Tensor[(300), float32] /* span=aten::linear_0.bias:0:0 */, %aten::linear_1.weight: Tensor[(300, 300), float32] /* span=aten::linear_1.weight:0:0 */, %aten::linear_1.bias: Tensor[(300), float32] /* span=aten::linear_1.bias:0:0 */, %aten::linear_2.weight: Tensor[(300, 300), float32] /* span=aten::linear_2.weight:0:0 */, %aten::linear_2.bias: Tensor[(300), float32] /* span=aten::linear_2.bias:0:0 */, %aten::linear_3.weight: Tensor[(10, 300), float32] /* span=aten::linear_3.weight:0:0 */, %aten::linear_3.bias: Tensor[(10), float32] /* span=aten::linear_3.bias:0:0 */) {
  %0 = nn.dense(%input0, %aten::linear_0.weight, units=None) /* span=aten::linear_0:0:0 */;
  %1 = nn.bias_add(%0, %aten::linear_0.bias, axis=-1) /* span=aten::linear_0:0:0 */;
  %2 = nn.relu(%1) /* span=aten::relu_0:0:0 */;
  %3 = nn.dense(%2, %aten::linear_1.weight, units=None) /* span=aten::linear_1:0:0 */;
  %4 = nn.bias_add(%3, %aten::linear_1.bias, axis=-1) /* span=aten::linear_1:0:0 */;
  %5 = nn.relu(%4) /* span=aten::relu_1:0:0 */;
  %6 = nn.dense(%5, %aten::linear_2.weight, units=None) /* span=aten::linear_2:0:0 */;
  %7 = nn.bias_add(%6, %aten::linear_2.bias, axis=-1) /* span=aten::linear_2:0:0 */;
  %8 = nn.relu(%7) /* span=aten::relu_2:0:0 */;
  %9 = nn.dense(%8, %aten::linear_3.weight, units=None) /* span=aten::linear_3:0:0 */;
  nn.bias_add(%9, %aten::linear_3.bias, axis=-1) /* span=aten::linear_3:0:0 */
}

Следующий шаг - компиляция модели без оптимизации слоев.

In [30]:
with tvm.transform.PassContext(opt_level=opt_level):
    lib = relay.build(mod, target=target, params=params)

После компиляции можно выполнить запуск вывода и измерение времени выполнения с использованием разработанной функции timeit_inference, а также проверку качества работы полносвязной нейронной сети после загрузки с помощью функции get_accuracy и сравнение полученной точности классификации с загруженным значением, которое получено на x86-64.

In [31]:
default_fcnn_predict, default_fcnn_times = timeit_inference(mod, lib, images)

default_fcnn_accuracy = get_accuracy(labels, default_fcnn_predict)
assert np.allclose(metric['fcnn'], default_fcnn_accuracy, rtol=1e-5)

default_fcnn_time = np.median(default_fcnn_times)
print(f'Медианное время работы не оптимизированной модели: {default_fcnn_time:.4f} мc')
Медианное время работы не оптимизированной модели: 1.8873 мc

7.2. Использование возможностей AutoTVM¶

Вызовем разработанную функцию get_autotvm_task для извлечения задач из графа вычислений для AutoTVM.

В данном случае следовало бы ожидать 8 задач, так как есть 4 слоя. Но задач 6: 3 с трансформацией данных и 3 без трансформации данных. Два слоя имеют идентичные параметры, поэтому данные задачи нет необходимости дублировать. Аналогичное поведение будет и у других методов оптимизации слоев.

In [32]:
tasks = get_autotvm_task(mod, target, params)
Извлечение задач

Номер задачи: 0
Информация о задаче: ('dense_nopack.x86', ('TENSOR', (1, 784), 'float32'), ('TENSOR', (300, 784), 'float32'), None, 'float32')

Номер задачи: 1
Информация о задаче: ('dense_pack.x86', ('TENSOR', (1, 784), 'float32'), ('TENSOR', (300, 784), 'float32'), None, 'float32')

Номер задачи: 2
Информация о задаче: ('dense_nopack.x86', ('TENSOR', (1, 300), 'float32'), ('TENSOR', (300, 300), 'float32'), None, 'float32')

Номер задачи: 3
Информация о задаче: ('dense_pack.x86', ('TENSOR', (1, 300), 'float32'), ('TENSOR', (300, 300), 'float32'), None, 'float32')

Номер задачи: 4
Информация о задаче: ('dense_nopack.x86', ('TENSOR', (1, 300), 'float32'), ('TENSOR', (10, 300), 'float32'), None, 'float32')

Номер задачи: 5
Информация о задаче: ('dense_pack.x86', ('TENSOR', (1, 300), 'float32'), ('TENSOR', (10, 300), 'float32'), None, 'float32')

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

In [33]:
log_file = 'autotvm/autotvm_fcnn.log'
n_trial = global_trial

tune_autotvm(tasks, n_trial, log_file)
[Task  2/ 6]  Current/Best:    1.72/   2.25 GFLOPS | Progress: (60/96) | 145.41 s Done.
[Task  2/ 6]  Current/Best:    2.43/   2.61 GFLOPS | Progress: (96/96) | 216.88 s Done.
[Task  4/ 6]  Current/Best:    1.10/   2.16 GFLOPS | Progress: (60/96) | 110.70 s Done.
[Task  4/ 6]  Current/Best:    1.12/   2.19 GFLOPS | Progress: (96/96) | 208.75 s Done.
[Task  5/ 6]  Current/Best:    0.05/   1.06 GFLOPS | Progress: (72/96) | 129.13 s Done.
[Task  6/ 6]  Current/Best:    0.42/   2.00 GFLOPS | Progress: (96/96) | 101.39 s Done.

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

In [34]:
with autotvm.apply_history_best(log_file):
    with tvm.transform.PassContext(opt_level=opt_level):
        lib = relay.build(mod, target=target, params=params)

На данном этапе можно выполнить измерение времени выполнения с использованием функции timeit_inference, проверку качества работы оптимизированной модели с помощью функции get_accuracy и сравнение точности классификации с рефенсным значением, которое было получено после запуска обучения модели.

In [35]:
autotvm_fcnn_predict, autotvm_fcnn_times = timeit_inference(mod, lib, images)

autotvm_fcnn_accuracy = get_accuracy(labels, autotvm_fcnn_predict)
assert np.allclose(metric['fcnn'], autotvm_fcnn_accuracy, rtol=1e-5)

autotvm_fcnn_time = np.median(autotvm_fcnn_times)
print(f'Медианное время работы после оптимизации слоев с помощью AutoTVM: {autotvm_fcnn_time:.4f} мc')
Медианное время работы после оптимизации слоев с помощью AutoTVM: 1.4901 мc

7.3. Применение MetaScheduler¶

Вызовем разработанную функцию get_ms_task, предварительно определив директорию work_dir для логирования результатов оптимизации.

В данном случае строка компиляции уже содержит информацию о числе потоков, поэтому модифицировать ее нет необходимости.

In [36]:
if is_x86():
    work_dir = "meta_schedule_fcnn"

    tasks, task_weights = get_ms_task(mod, target, params, opt_level, work_dir)

Далее выполним запуск оптимизации с помощью MetaScheduler посредством вызова функции tune_ms, установив число экспериментов при оптимизации равным N * len(tasks).

In [37]:
n_trial_per_task = global_trial

if is_x86():
    tune_ms(tasks, task_weights, work_dir, n_trial_per_task * len(tasks))

После оптимизации можно скомпилировать нейронную с учетом построенных оптимизаций с помощью интерфейса MetaScheduler ms.relay_integration.compile_relay.

In [38]:
if is_x86():
    
    database = ms.database.JSONDatabase(
        f"{work_dir}/database_workload.json",
        f"{work_dir}/database_tuning_record.json",
        allow_missing=False
    )

    lib = ms.relay_integration.compile_relay(
        database, mod, target, params,
        opt_level=opt_level,
    )

В завершении измерим время вывода с использованием функции timeit_inference, определим качество работы модели с помощью функции get_accuracy и выполним проверку корректности работы оптимизированной модели, сравнив полученное значение показателя точности с референсным.

In [39]:
if is_x86():
    ms_fcnn_predict, ms_fcnn_times = timeit_inference(mod, lib, images)

    ms_fcnn_accuracy = get_accuracy(labels, ms_fcnn_predict)
    assert np.allclose(metric['fcnn'], ms_fcnn_accuracy, rtol=1e-5)

    ms_fcnn_time = np.median(ms_fcnn_times)
    print(f'Медианное время работы после оптимизации слоев с помощью MetaScheduler: {ms_fcnn_time:.4f} мc')

7.4. Анализ результатов¶

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

In [40]:
fig, ax = plt.subplots()

name = ['Без оптимизации\nслоев', 'AutoTVM', 'MetaScheduler']
times = [default_fcnn_time, autotvm_fcnn_time, ms_fcnn_time]

bars = ax.bar(name, times, label=name, color=bar_colors)
ax.set_title('Среднее время\nвыполнения (мс)', fontsize=18)

for bar, n, t in zip(bars, name, times):
    h = bar.get_height()
    if n == 'Без оптимизации\nслоев': h = h / 2
    if h != 0:
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            h,
            f'{round(t, 4)} с',
            ha='center',
            va='bottom',
            fontsize=15,
        )

ax.xaxis.label.set_size(40)
ax.set_title('Среднее время\nвыполнения (с)', fontsize=18)
plt.grid()
No description has been provided for this image

Вывод: оптимизация значительно ускоряет время работы сети.

8. Запуск и оптимизация сверточной нейронной сети¶

8.1. Компиляция и запуск модели¶

Вначале необходимо выполнить загрузку модели сверточной нейронной сети.

In [41]:
default_cnn_time, autotvm_cnn_time, autoscheduler_cnn_time, ms_cnn_time = 0, 0, 0, 0

mod, params = load_model('model/cnn.json', 'model/cnn.params')
print(mod['main'])
fn (%input0: Tensor[(1, 1, 28, 28), float32] /* span=aten::_convolution_0.input0:0:0 */, %aten::_convolution_0.weight: Tensor[(64, 1, 3, 3), float32] /* span=aten::_convolution_0.weight:0:0 */, %aten::_convolution_0.bias: Tensor[(64), float32] /* span=aten::_convolution_0.bias:0:0 */, %aten::linear_0.weight: Tensor[(10, 12544), float32] /* span=aten::linear_0.weight:0:0 */, %aten::linear_0.bias: Tensor[(10), float32] /* span=aten::linear_0.bias:0:0 */) {
  %0 = nn.conv2d(%input0, %aten::_convolution_0.weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* span=aten::_convolution_0:0:0 */;
  %1 = nn.bias_add(%0, %aten::_convolution_0.bias) /* span=aten::_convolution_0:0:0 */;
  %2 = nn.relu(%1) /* span=aten::relu_0:0:0 */;
  %3 = nn.max_pool2d(%2, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* span=aten::max_pool2d_0:0:0 */;
  %4 = reshape(%3, newshape=[1, -1]) /* span=aten::view_0:0:0 */;
  %5 = nn.dense(%4, %aten::linear_0.weight, units=None) /* span=aten::linear_0:0:0 */;
  nn.bias_add(%5, %aten::linear_0.bias, axis=-1) /* span=aten::linear_0:0:0 */
}

Следующий шаг - компиляция модели без оптимизации слоев.

In [42]:
with tvm.transform.PassContext(opt_level=opt_level):
    lib = relay.build(mod, target=target, params=params)

После компиляции можно выполнить запуск вывода и измерение времени выполнения с использованием разработанной функции timeit_inference, а также проверку качества работы полносвязной нейронной сети после загрузки с помощью функции get_accuracy и сравнение полученной точности классификации с загруженным значением, которое получено на x86-64.

In [43]:
default_cnn_predict, default_cnn_times = timeit_inference(mod, lib, images)

default_cnn_accuracy = get_accuracy(labels, default_cnn_predict)
assert np.allclose(metric['cnn'], default_cnn_accuracy, rtol=1e-5)

default_cnn_time = np.median(default_cnn_times)
print(f'Медианное время работы не оптимизированной модели: {default_cnn_time:.4f} мc')
Медианное время работы не оптимизированной модели: 0.8829 мc

8.2. Использование возможностей AutoTVM¶

Вызовем разработанную функцию get_autotvm_task для извлечения задач из графа вычислений для AutoTVM.

В данном случае к задачам с полносвязным слоем добавляется задача со сверточным слоем.

In [44]:
if is_x86():
    tasks = get_autotvm_task(mod, target, params)

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

In [45]:
log_file = 'autotvm/autotvm_cnn.log'
n_trial = global_trial

if is_x86():
    tune_autotvm(tasks, n_trial, log_file)

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

In [46]:
if is_x86():
    with autotvm.apply_history_best(log_file):
        with tvm.transform.PassContext(opt_level=opt_level):
            lib = relay.build(mod, target=target, params=params)

На данном этапе можно выполнить измерение времени выполнения с использованием функции timeit_inference, проверку качества работы оптимизированной модели с помощью функции get_accuracy и сравнение точности классификации с рефенсным значением, которое было получено после запуска обучения модели.

In [47]:
if is_x86():
    autotvm_cnn_predict, autotvm_cnn_times = timeit_inference(mod, lib, images)

    autotvm_cnn_accuracy = get_accuracy(labels, autotvm_cnn_predict)
    assert np.allclose(metric['cnn'], autotvm_cnn_accuracy, rtol=1e-5)

    autotvm_cnn_time = np.median(autotvm_cnn_times)
    print(f'Медианное время работы после оптимизации слоев с помощью AutoTVM: {autotvm_cnn_time:.4f} мc')

8.3. Использование Auto-scheduler¶

Вызовем разработанную функцию get_auto_scheduler_task для извлечения задач из графа вычислений для AutoTVM.

In [48]:
if is_x86():
    tasks, task_weights = get_auto_scheduler_task(mod, target, params, opt_level)

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

In [49]:
os.makedirs('auto_schedule/', exist_ok=True)
log_file = 'auto_schedule/auto-schedule_cnn.log'
n_trial_per_task = global_trial

if is_x86():
    tune_auto_scheduler(tasks, task_weights, log_file, n_trial_per_task * len(tasks))

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

In [50]:
if is_x86():
    with auto_scheduler.ApplyHistoryBest(log_file):
        with tvm.transform.PassContext(
            opt_level=opt_level, config={"relay.backend.use_auto_scheduler": True},
        ):
            lib = relay.build(mod, target=target, params=params)

На данном этапе можно выполнить измерение времени выполнения с использованием функции timeit_inference, проверку качества работы оптимизированной модели с помощью функции get_accuracy и сравнение точности классификации с рефенсным значением, которое было получено после запуска обучения модели.

In [51]:
if is_x86():
    autoscheduler_cnn_predict, autoscheduler_cnn_times = timeit_inference(mod, lib, images)

    autoscheduler_cnn_accuracy = get_accuracy(labels, autoscheduler_cnn_predict)
    assert np.allclose(metric['cnn'], autoscheduler_cnn_accuracy, rtol=1e-5)

    autoscheduler_cnn_time = np.median(autoscheduler_cnn_times)
    print(f'Медианное время работы после оптимизации слоев с помощью Auto-scheduler: {autoscheduler_cnn_time:.4f} мc')

8.4. Применение MetaScheduler¶

Вызовем разработанную функцию get_ms_task, предварительно определив директорию work_dir для логирования результатов оптимизации.

В данном случае строка компиляции уже содержит информацию о числе потоков, поэтому модифицировать ее нет необходимости.

In [52]:
if is_x86():
    work_dir = "meta_schedule_cnn"

    tasks, task_weights = get_ms_task(mod, target, params, opt_level, work_dir)

Далее выполним запуск оптимизации с помощью MetaScheduler посредством вызова функции tune_ms, установив число экспериментов при оптимизации равным N * len(tasks).

In [53]:
n_trial_per_task = global_trial

if is_x86():
    tune_ms(tasks, task_weights, work_dir, n_trial_per_task * len(tasks))

После оптимизации можно скомпилировать нейронную с учетом построенных оптимизаций с помощью интерфейса MetaScheduler ms.relay_integration.compile_relay.

In [54]:
if is_x86():
    
    database = ms.database.JSONDatabase(
        f"{work_dir}/database_workload.json",
        f"{work_dir}/database_tuning_record.json",
        allow_missing=False
    )

    lib = ms.relay_integration.compile_relay(
        database, mod, target, params,
        opt_level=opt_level,
    )

В завершении измерим время вывода с использованием функции timeit_inference, определим качество работы модели с помощью функции get_accuracy и выполним проверку корректности работы оптимизированной модели, сравнив полученное значение показателя точности с референсным.

In [55]:
if is_x86():
    ms_cnn_predict, ms_cnn_times = timeit_inference(mod, lib, images)

    ms_cnn_accuracy = get_accuracy(labels, ms_cnn_predict)
    assert np.allclose(metric['cnn'], ms_cnn_accuracy, rtol=1e-5)

    ms_cnn_time = np.median(ms_cnn_times)
    print(f'Медианное время работы после оптимизации слоев с помощью MetaScheduler: {ms_cnn_time:.4f} мc')

8.5. Анализ результатов¶

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

In [56]:
fig, ax = plt.subplots()

name = ['Без оптимизации\nслоев', 'AutoTVM', 'Auto-scheduler', 'MetaScheduler']
times = [default_cnn_time, autotvm_cnn_time, autoscheduler_cnn_time, ms_cnn_time]

bars = ax.bar(name, times, label=name, color=bar_colors)
ax.set_title('Среднее время\nвыполнения (мс)', fontsize=18)

for bar, n, t in zip(bars, name, times):
    h = bar.get_height()
    if n == 'Без оптимизации\nслоев': h = h / 2
    if h != 0:
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            h,
            f'{round(t, 4)} с',
            ha='center',
            va='bottom',
            fontsize=15,
        )

ax.xaxis.label.set_size(40)
ax.set_title('Среднее время\nвыполнения (с)', fontsize=18)
plt.grid()
No description has been provided for this image

Вывод: оптимизация значительно ускоряет время работы сети.