Экспорт нейросети в формат ONNX

Создаём нейросеть и сохраняем ее в формате ONNX для дальнейшего использования.

Описание задачи

Чтобы развернуть готовую нейросеть в целевом окружении или даже встроить в код для микроконтроллера, распространённым решением является перевод нейросети в стандартный, независимый от фреймворков формат ONNX (Open Neural Network eXchange).

Самой распространённой библиотека для создания и обучения нейросетей в Julia называется Flux. Она позволяет создавать полносвязанные, сверточные, рекуррентные нейросети и графы других топологий, писать собственные модели нейронов и слоёв, а обучение и прямой расчет производятся в скомпилированной, бинарной форме, по скорости не уступающей языку Си.

Функции экспорта Flux-нейросетей в ONNX встречаются в нескольких библиотеках, но на данный момент наиболее успешной является инструмент экспорта, входящий в библиотеку NaiveNAXflux. Эта библиотека содержит инструменты для взаимосвязи библиотеки Flux с библиотекой для поиска оптимальных топологий нейросетей (NAS - Neural Architecture Search). NaiveNAS берет на себя работу с весами и изменение топологии нейросети, а сопряженные с ним пакеты, вроде NaiveGAflux, управляют генетической оптимизацией.

Мы также установим библиотеку ONNXRunTime, которая позволит нам запускать и тестировать сохранённую нейросеть.

Сохранение нейросети в ONNX

Установим нужные библиотеки:

Pkg.add(["Flux", "ONNXNaiveNASflux", "ONNXRunTime", "ONNXLowLevel"])

После их подключения мы сможем экспортировать в ONNX простую полносвязанную сеть:

using Flux
using ONNXNaiveNASflux

# Определим модель
model = Chain(
    Dense(784 => 32, relu),
    Dense(32 => 10)
)

# Экспортируем в ONNX при помощи save()
ONNXNaiveNASflux.save("fc_model.onnx", model, (784, 1));

Последним аргументом этой команды мы указываем размерность входных данных в формате количество признаков, размер батча (features, batch_size).

Количество признаков в описании объекта можно получить из свойств нейросети Flux, но размер батча во Flux не зафиксирован.

first_layer = model[1]
input_size = size(first_layer.weight, 2)

784

Выполнение ранее сохраненной нейросети

Запустим сохранённую нейросеть (веса в ней инициализированы случайными числами, поэтому на выходе мы получим шум):

import ONNXRunTime as ORT
fc_model = ORT.load_inference("fc_model.onnx")
fc_model

InferenceSession
    input_names:        ["data_0"]
    output_names:       ["dense_1"]
    execution_provider: :cpu

Теперь используем названия весов чтобы задать входы и получить выходы из сохраненной ONNX нейросети:

result = []

for x in eachcol(rand(Float32, 784, 20))
    input_data = Dict("data_0" => reshape(x, 1, :))
    outputs = fc_model(input_data)
    append!(result, [outputs["dense_1"]])
end

gr()
plot(vcat(result...), title="Прогноз нейросети для случайных объектов")

Экспорт других топологий

Вот экспорт основных слоёв рекуррентной нейросети с одним RNN слоем.

using Flux
using ONNXNaiveNASflux

# Эта модель принимает 5 переменных, производит их обработку при помощи RNN и выдает 2 значения
rnn_model = Chain(
    RNN(5 => 10),      # Входы: 5 признаков, Внутренний слой: 10 признаков
    Dense(10 => 2)     # Выходной слой переводит состояние скрытого слоя в два класса/значения
)

# Определим размерность входного тензора:
# Переменных = 5, Размер батча = 1, Длина входной последовательности = 100
input_shape = (5, 1, 100)

ONNXNaiveNASflux.save("rnn_model.onnx", rnn_model, input_shape);

Запустим рекуррентную сеть из файла ONNX:

import ONNXRunTime as ORT
onnx_rnn_model = ORT.load_inference("rnn_model.onnx")
onnx_rnn_model

InferenceSession
    input_names:        ["data_0"]
    output_names:       ["dense_0"]
    execution_provider: :cpu

С этой информацией мы можем запустить рекуррентную нейросеть и прогнозировать выходные признаки для каждого отсчета входного временного ряда:

plot()
for x in 1:20
     # Генерируем тестовую последовательность : 5 признаков на 100 отсчетов
     raw_data = rand(Float32, 5, 100)

     # Добавляем размерность с учетом батча (стандарт ONNX)
     input_tensor = reshape(raw_data', 100, 1, 5) 
     outputs = onnx_rnn_model(Dict("data_0" => input_tensor))
     onnx_output = outputs["dense_0"]

     # Пересобираем данные в матрицу (100 строк, 2 колонки)
     plot_data = reshape(onnx_output, 100, 2)

     gr()
     plot!(plot_data, title="Прогноз RNN для 20 случайных запусков",
         label = (x==1 ? ["Выход 1" "Выход 2"] : false),
         line = (x==1 ? (5,:solid) : (1,:dot)),
         c = [:red :black],
         xlabel="Временной шаг (отсчет)")
end
plot!()

Сверточные сети

Если соблюдать порядок аргументов, то несложно выполнить то же самое для свёрточных нейросетей:

using Flux
using ONNXNaiveNASflux

# Простая сверточная сеть
# Вход: 28x28 (1 канал), Выход: 10 классов (например, цифры 0-9)
cnn_model = Chain(
    Conv((3, 3), 1 => 4, relu), # Свертка: фильтр 3x3, 1 входной канал, 4 выходных
    Flux.flatten,               # Выпрямление тензора в вектор
    Dense(26 * 26 * 4 => 10)    # Полносвязный слой для классификации
)

# Задаем размерность в формате Flux/Julia: (Width=28, Height=28, Channels=1, Batch=1)
input_shape_cnn = (28, 28, 1, 1)

# Экспортируем
ONNXNaiveNASflux.save("cnn_model.onnx", cnn_model, input_shape_cnn);

И запустим сохранённую в ONNX свёрточную нейросеть:

import ONNXRunTime as ORT

# Загружаем модель
onnx_cnn_model = ORT.load_inference("cnn_model.onnx")

# 1. Генерируем случайную "картинку" 28x28 с 1 каналом в формате Julia (W, H, C)
# В реальной задаче здесь будет реальное изображение
raw_image = rand(Float32, 28, 28, 1) 

# 2. Переставляем оси для ONNX
# Из (W=28, H=28, C=1) делаем (C=1, H=28, W=28)
permuted_image = permutedims(raw_image, (3, 2, 1))

# Добавляем размерность батча на первую позицию -> получаем (Batch=1, C=1, H=28, W=28)
input_tensor = reshape(permuted_image, 1, 1, 28, 28)

# 3. Передаем в модель (имя входа для первого слоя обычно "data_0")
input_data = Dict("data_0" => input_tensor)
outputs = onnx_cnn_model(input_data)

# 4. Извлекаем выходные данные
# На выходе будет тензор с логитами (вероятностями) для 10 классов
onnx_output = outputs["dense_0"]

# 5. Принудительно приводим к одномерному вектору (10 элементов) для графика
# Это защитит нас от любой структуры батчей на выходе (будь то 1x10 или 10x1)
plot_data = reshape(onnx_output, :)

# Визуализация прогноза по 10 классам
gr()
bar(plot_data, title="Прогноз CNN для случайного изображения", 
    xticks=1:10, xlabel="Классы (0-9)", ylabel="Уверенность сети", legend=false)

Заключение

Мы разобрались с тем, как сохранить в ONNX нейросети трех основных видов. Не все конструкции из Flux экспортируются надежным образом, или вообще имеют поддержку в ONNXNaiveNASflux. Могут потребоваться эксперименты и упрощение вашей топологии.