Проект Google Coral недавно вышел из стадии бета-тестирования. Согласно тестам, устройства Coral обеспечивают отличное ускорение вывода нейронных сетей для домашних мастеров. Эти устройства основаны на специализированном блоке тензорной обработки ASIC (Edge TPU), работать с которым оказалось довольно сложно, но принудительные ограничения и причуды вознаграждаются. Мне не терпелось исследовать внутренние особенности взаимодействия между TensorFlow и Edge TPU и взломать оба, чтобы делать крутые, нестандартные, безумные вещи.
Предполагается, что вы знакомы с основами TensorFlow и Edge TPU. Официальная документация хороша, поэтому достаточно просмотреть модели TensorFlow на Edge TPU и Edge TPU Compiler, чтобы продолжить. Для повторения моих экспериментов требуется Ubuntu Linux и внешний Coral USB Accelerator.
Во-первых, программное обеспечение Edge TPU не является полностью открытым исходным кодом. Самые «вкусные» биты, edgetpu_compiler исполняемый файл иlibedgetpu.so разделяемая библиотека - проприетарные. Этот факт увеличивает потенциальную сложность взлома, но также делает его более увлекательным! Например, единственный способ узнать, какие API предоставляет libedgetpu.so, - это сбросить экспортированные символы с objdump:
$ objdump -TCj .text /usr/lib/x86_64-linux-gnu/libedgetpu.so.1 /usr/lib/x86_64-linux-gnu/libedgetpu.so.1: file format elf64-x86-64 DYNAMIC SYMBOL TABLE: 000000000006baa0 g DF .text 000000000000000d VER_1.0 edgetpu::RegisterCustomOp() 0000000000072b40 g DF .text 000000000000001f VER_1.0 edgetpu::EdgeTpuContext::~EdgeTpuContext() 0000000000072ad0 g DF .text 0000000000000006 VER_1.0 edgetpu::EdgeTpuContext::~EdgeTpuContext() 0000000000072ad0 g DF .text 0000000000000006 VER_1.0 edgetpu::EdgeTpuContext::~EdgeTpuContext() 000000000006dc10 g DF .text 000000000000000a VER_1.0 tflite_plugin_destroy_delegate 000000000006be50 g DF .text 00000000000001dd VER_1.0 edgetpu_list_devices 000000000006bb80 g DF .text 0000000000000107 VER_1.0 edgetpu_version 000000000006bab0 g DF .text 000000000000000a VER_1.0 edgetpu::EdgeTpuManager::GetSingleton() 000000000006d090 g DF .text 0000000000000b7c VER_1.0 tflite_plugin_create_delegate 000000000006bb20 g DF .text 0000000000000012 VER_1.0 edgetpu_free_devices 000000000006bac0 g DF .text 000000000000005e VER_1.0 edgetpu::operator<<(std::ostream&, edgetpu::DeviceType) 000000000006c030 g DF .text 0000000000000c1a VER_1.0 edgetpu_create_delegate 000000000006bb40 g DF .text 000000000000000a VER_1.0 edgetpu_free_delegate 000000000006bb50 g DF .text 0000000000000024 VER_1.0 edgetpu_verbosity
Этот вывод явно подразумевает, что Edge TPU API прибит к TensorFlow Lite длинными толстыми гвоздями. По-другому пользоваться устройством просто невозможно. Если вы ожидали увидеть API нижнего уровня вроде «умножьте эту матрицу на этот вектор на Edge TPU» - неудача.
Итак, давайте быстро резюмируем, как работает TensorFlow Lite API для Edge TPU:
- Создайте граф вычислений с помощью обычного TensorFlow. Например, обучите глубокую нейронную сеть.
- Преобразуйте его в формат TensorFlow Lite, который представляет собой плоские буферы вместо protobuf и с другой схемой. Новый граф должен быть особенным, чтобы подружиться с Edge TPU. Примечательно, что содержащиеся операции (ops) должны быть квантованы в
uint8, потому что Edge TPU может работать только с байтами без знака. - Скрестите пальцы и преобразуйте его еще раз, на этот раз с
edgetpu_compiler. Базовый формат остается прежним, но поддерживаемые операции объединяются и компилируются в один волшебный блок Edge TPU. - Убедитесь, что устройство Coral подключено, создайте новый интерпретатор TensorFlow Lite с делегатом операций Edge TPU и вызовите.
Таким образом, невозможно выполнить произвольный расчет на Edge TPU без вызова внешних программ и записи и чтения файлов на ходу.
Важнейшей деталью этой многоступенчатой процедуры является список операций, которые можно скомпилировать для Edge TPU. TensorFlow Lite не поддерживает все удары и свистки своего старшего брата, а Edge TPU поддерживает только часть того, что осталось. Например, нет умножения матриц (сюрприз! Я эмпирически проверил, что «выходной тензор одномерный» в FullyConnected). Эти ограничения делают что-либо, кроме вывода сверточной и полностью связанной нейронной сети на Edge TPU, трудным, очень сложным. Но не невозможно. Вот где пост становится забавным.
Размытие в движении по краю TPU
Эффект размытия движения является результатом двумерной свертки изображения с ядром «радиуса».
В терминах TensorFlow эта операция называется DepthwiseConv2d, она широко используется в глубоких сверточных нейронных сетях и поддерживается Edge TPU. Пиксели изображения могут быть представлены в формате RGB, по одному байту на канал - именно то, что нужно Edge TPU. Давайте пройдемся по всем ямкам и опасностям и посмотрим, насколько быстро выполняется фильтрация изображений размытия движения в Edge TPU с Python!
0 → tf.функция
Забудьте о существовании TensorFlow Lite и Edge TPU в этом разделе. Познакомимся с основной логикой. Следующий код создает сверточное ядро. dim - размер, а angle - угол движения на плоскости в радианах.
matplotlib всегда удобен, когда дело касается грязных визуальных эффектов.
Следующим шагом является проверка нашего эффекта размытия движения с помощью обычного TensorFlow 2.0. Мы используем tf.nn.depthwise_conv2d для вычисления двумерной свертки изображения с помощью нашего ядра. Все шаги равны 1, так что размеры изображения не меняются.
В Jupyter можно быстро измерить эффективность размытия при движении с помощью %timeit motion_blur(images). Это дает примерно 5,30 с ± 0,09 на моем процессоре Intel i7–8565U 4x2 (HT).
tf.function → tflite
Теперь, когда мы уверены, что общий подход работает, пришло время перенести его на TensorFlow Lite.
Мы должны указать входную сигнатуру tf.functionin create_motion_blur_func, потому что TensorFlow Lite в настоящее время не допускает переменных форм, кроме первого «пакетного» измерения. Следовательно, наше размытие в движении может работать только с изображениями одинакового размера.
create_motion_blur_func_lite - это оболочка для create_motion_blur_func, конвертирующая последний в TensorFlow Lite. generate_lite_model инициализирует tf.lite.TFLiteConverter из графа вычислений, принадлежащего tf.function - нашему алгоритму размытия движения - и записывает результат преобразования на диск. create_func_lite загружает его обратно, устанавливает новый tf.lite.Interpreter и возвращает закрытие вызова.
Согласно %timeit, новая реализация быстрее: 3,50 с ± 0,24 против 5,3 с. Этот прирост производительности удивителен, потому что, согласно системному монитору, при выполнении используется только одно из восьми ядер ЦП. Мы можем визуализировать получившуюся модель .tflite с помощью netron:
tflite → Edge TPU
Наконец, нам нужно перейти с обычного TensorFlow Lite на Edge TPU. Этот шаг, безусловно, самый хитрый и сложный. Мы продолжим работу над существующим кодом, добавляя по одной функции за раз.
Edge TPU требует uint8 тип данных операции (dtype) вместо float32. К сожалению, мы не можем заставить tf.nn.depthwise_conv2d работать напрямую с uint8: поддерживаются только float64, float32, bfloat16 и float16. Следовательно, мы должны прибегать к «квантованию после обучения», что означает подделку d-типов и добавление свойств квантования ко всем операциям. gen_input_samples имитирует диапазон значений пикселей от 0 до 255, именно так квантование параметризуется в TensorFlow Lite. Далее мы вызываем edgetpu_compiler в квантованной модели, чтобы заменить операцию двумерной свертки оптимизированным кодом для Edge TPU. tf.lite.Interpreter необходимо дополнить experimental_delegates=[load_delegate("libedgetpu.so.1.0")], чтобы он знал, что делать с оптимизированной операцией Edge TPU.
В идеальном мире, где edgetpu_compiler поддерживает TensorFlow 2.0, приведенный выше код должен работать. Запустим код и посмотрим.
Edge TPU Compiler version 2.0.267685300 Model compiled successfully in 231 ms. Input model: motion_bluredgetpu_compiler1920_1058libedgetpu.so25_1.57.tflite Input size: 3.03KiB Output model: motion_bluredgetpu_compiler1920_1058libedgetpu.so25_1.57_edgetpu.tflite Output size: 296.85KiB On-chip memory available for caching model parameters: 1.73MiB On-chip memory used for caching model parameters: 10.00KiB Off-chip memory used for streaming uncached model parameters: 0.00B Number of Edge TPU subgraphs: 1 Total number of operations: 3 Operation log: motion_bluredgetpu_compiler1920_1058libedgetpu.so25_1.57_edgetpu.log Model successfully compiled but not all operations are supported by the Edge TPU. A percentage of the model will instead run on the CPU, which is slower. If possible, consider updating your model to use only operations supported by the Edge TPU. For details, visit g.co/coral/model-reqs. Number of operations that will run on Edge TPU: 1 Number of operations that will run on CPU: 2 Operator Count Status DEPTHWISE_CONV_2D 1 Mapped to Edge TPU DEQUANTIZE 1 Operation is working on an unsupported data type QUANTIZE 1 Operation is otherwise supported, but not mapped due to some unspecified limitation
DEPTHWISE_CONV_2D успешно компилируется, однако есть странные DEQUANTIZE и QUANTIZE операции, которые этого не делают. Это артефакты TensorFlow 2.0, которые не поддерживаются компилятором и возникли из принудительного float32 dtype в подписи motion_blur_func. netron визуализация должна все прояснить.
Таким образом, мы должны выполнить дублирующую работу четыре раза:
- Переключите значения пикселей с
uint8наfloat32, передайте их механизму TensorFlow Lite. - TensorFlow Lite выполняет
QUANTIZEи снова переключается наuint8 - Вычислив свертку, мы возвращаемся к
float32вDEQUANTIZE. - TensorFlow Lite возвращает управление вызывающей стороне, и мы конвертируем в
uint8, чтобы сохранить изображение.
Если отбросить QUANTIZE и DEQUANTIZE из исходного .tflite, модель снова станет великолепной. К сожалению, нет простого способа добиться этого. Официального API для управления моделями TensorFlow Lite просто не существует. Нам нужно копать глубже.
Копать глубже
Я упоминал, что формат модели TensorFlow Lite - это плоские буферы. Это относительно новый универсальный формат сериализации, который, я уверен, конкурирует с protobuf внутри Google. API Python flatbuffers не позволяет изменять существующие файлы, облом. Нам повезло, что плоский буфер демонстрирует два разных представления объекта: двоичное и JSON, и поддерживает преобразование между ними без потерь. Существуют команды flatc -j и flatc -b для преобразования соответственно из .tflite в JSON и обратно. Мы собираемся применить их, чтобы убрать избыточные операции из модели, предоставляемой схема является общедоступной. Фактически, это метод, которым разработчики TensorFlow Lite сами себя используют для обновления моделей .tflite.
Код показывает, что исходная модель .tflite имела неправильные dtypes: int8 вместо uint8. TensorFlow 2.0 пытается применить многоканальное квантование без запроса; Edge TPU не поддерживает многоканальное квантование, и это тоже нужно исправить.
Вторая попытка более удачна:
Edge TPU Compiler version 2.0.267685300 Model compiled successfully in 171 ms. Input model: motion_bluredgetpu_compiler1920_1058libedgetpu.so25_1.57.tflite Input size: 2.71KiB Output model: motion_bluredgetpu_compiler1920_1058libedgetpu.so25_1.57_edgetpu.tflite Output size: 296.56KiB On-chip memory available for caching model parameters: 1.73MiB On-chip memory used for caching model parameters: 10.00KiB Off-chip memory used for streaming uncached model parameters: 0.00B Number of Edge TPU subgraphs: 1 Total number of operations: 1 Operation log: motion_bluredgetpu_compiler1920_1058libedgetpu.so25_1.57_edgetpu.log Operator Count Status DEPTHWISE_CONV_2D 1 Mapped to Edge TPU
Теперь обещанный тест. Я установил libedgetpu-max, который не ограничивает рабочую частоту. Мои результаты составляют 5,00 с ± 0,25 и 0,262 с ± 0,001 для исходной версии и версии Edge TPU соответственно. Edge TPU в 10–20 раз быстрее, чем самая быстрая float32 реализация на моем процессоре! Это сравнение, конечно, нечестно, поскольку для запуска исходного .tflite используется только одно ядро ЦП, и я не могу изменить его в Python (это выглядит возможным в C ++). Я ожидаю, что реальное ускорение производительности составит 2–4 раза. Кроме того, правильная реализация CPU vectorizeduint8 должна быть в 4 раза быстрее, чем float32 - например, Pillow-simd. Таким образом, Edge TPU несправедливого превосходства. С другой стороны, устройство Coral потребляет как минимум в 20 раз меньше энергии.
Изображение, созданное на Edge TPU, выглядит идентично наземному, но не побайтно из-за потери точности.
Изменение размера Ланцоша на Edge TPU
Размытие изображений - это весело, но есть еще более практическое применение tf.nn.depthwise_conv2d. Пожалуй, самый качественный метод изменения размера изображения также основан на свертках. Свертка действует как фильтр нижних частот перед передискретизацией пикселей. Существуют различные ядра с усреднением пикселей, и, пожалуй, наиболее известным является Lanczos. Наш план будет заключаться в определении нового ядра Lanczos в TensorFlow и создании модели Edge TPU для сжатия изображений в 2 раза.
Мы повторно используем функции, которые мы разработали ранее: create_func_edgetpu и generate_edgetpu_model. Самое интересное место в коде - это то, как мы используем сверточное сканирование ядра для реализации субдискретизации пикселей в той же операции. Посмотрим на результаты.
Показывая попугая
Ой. Обычный белый прямоугольник. Что-то пошло не так. Причина кроется в том, как выглядит ядро Ланцоша. Давайте визуализируем это так же, как мы делали это в разделе размытия в движении.
Вы видите «-0,006» на цветной шкале? Правильно, ядро Ланцоша содержит отрицательные значения. Как вы помните, встроенное квантование меняет float32 на int8, а моя «постобработка JSON» устанавливает dtypes на uint8. Следовательно, ядро применено неправильно, и мы терпим крушение с массовыми переполнениями. Мы должны поднять нулевую точку квантования с нуля и обновить все веса ядра.
Маскировка переполнений
Наше ядро не идеально: из-за ошибок округления и общей неточности крошечная часть пиксельных каналов оказалась больше 256. Вот почему бирюзовые артефакты проявлялись возле глаз попугая там, где погода белая. Точно так же самые темные области могут стать отрицательными. Нам нужно отсечение значений.
Результаты тестов: 150 мс ± 2 для _65 _ / _ 66_, 165 мс ± 3 для tflite / float32, 220 мс ± 3 для tflite / uint8 и 47,8 мс ± 0,1 для Edge TPU. То есть Edge TPU в 3 раза быстрее. Изменение размера того же изображения с помощью подушки-simd занимает 4,9 мс ± 0,1, поэтому Edge TPU примерно в 10 раз медленнее, чем одна из самых быстрых реализаций изменения размера Lanczos на моем процессоре. Я использовал следующий код:
from PIL import Image
img = Image.open("parrot.jpg")
%timeit img.resize((529, 960), Image.LANCZOS)
Я не знаю точного размера ядра, используемого в подушке-simd, вероятно, 5x5 вместо наших 11x11. Изменение размера ядра на 5x5 дает 40 мс на Edge TPU. Такое незначительное изменение времени является показателем того, что вся мощь Edge TPU раскрывается с большими размерами ядра.
Итак, я тестировал с различными размерами ядра. edgetpu_compiler аварийно завершает работу при размере ядра больше или равном 28.
Очевидно, что на стороне TensorFlow / Python существует огромный постоянный фактор, который предшествует преимуществу Edge TPU. Ядра 25x25 дают в 25 раз больше вычислений, но затраченное время увеличивается только вдвое. Этот факт согласуется с заявлением документации Coral о том, чтобы максимально избегать связи CPU-Edge TPU.
Выводы
- Использование Edge TPU возможно только с TensorFlow Lite.
- На Edge TPU сложно запрограммировать что-либо, кроме сверточного вывода NN. Но возможно. Например, универсальная фильтрация изображений с использованием
tf.nn.depthwise_conv2d. - Edge TPU не любит TensorFlow 2.0. Лучше придерживаться того, что написано в документации, и использовать 1.x. Еще можно взломать 2.0.
- Модели TensorFlow Lite можно свободно редактировать и взламывать в представлении JSON.
- Преобразования на Edge TPU будут привязаны к определенному размеру изображения. Однако можно предварительно сгенерировать модели.
- Размытие в движении или другие фильтры свертки можно рассчитать на Edge TPU с незначительной визуальной разницей. Производительность не уступает современному процессору Intel 4x2HT или лучше при больших размерах ядра.
- Изменение размера Lanczos на Edge TPU выполняется быстро, но все же в 10 раз медленнее, чем отличная реализация векторизованного процессора.
- Продемонстрированные приемы обработки изображений должны лучше всего работать как часть конвейера увеличения, включенного в модель CNN.
- Edge TPU демонстрирует полную мощность, когда хост io не доминирует.
Я рекомендую читателю попробовать другие ядра свертки на Edge TPU. Различные виды размытия, обнаружения линий, краев, повышения резкости и т. Д. И т. Д. Выполните поиск по запросу «примеры свертки изображений».
