Что внутри у Google Coral Edge TPU: тесты на скорость и разбор устройства

В 2019 году Google, наконец, выпустила оборудование с TPU под брендом Coral, которое можно купить. Однако не такие мощные облачные сети для обучения типа BigGAN со своими 100 петафлопс/с, и даже не такие, как самые дешёвые 180 TFlop/s v2 TPU, которые можно арендовать за $4,5 в час. Эти TPU-устройства должны работать «на краю», то есть, обеспечивать возможность внедрения решений с глубинным обучением для малых устройств без доступа к интернету. Насколько же хороши будут те TPU, которые мы сможем купить?

В спецификации текущих продуктов их скорость указана, как 4 TOPs (тераопераций в секунду). Это не стандартные устройства с поддержкой операций с 32-битными числами с плавающей запятой, поскольку tensorflow-lite, работающий на этих устройствах, использует 8-битную арифметику с фиксированной запятой. Также вы не найдёте у них такой полосы пропускания для памяти, как на облачных TPU или GPU. «Краевые» TPU подсоединяются к устройствам через USB 3.0 или по шине mPCIe 2-го поколения, поэтому скорость ограничится 500-640 МБ/с. Конечно, поскольку на таком устройстве есть всего 8 МБ SRAM, на передачу данных на или с устройства уйдёт не более 16 мс, а у модели, которую мы тестировали для написания этой статьи, на это ушло всего 10 мс. Вы не будете страдать от нехватки памяти, поскольку эти устройства не предназначены для обучения с нуля. Однако на этом устройстве можно обучить последний слой модели. Основной смысл использования краевых TPU – построение логических заключений, и именно так я буду проверять скорость данного устройства.

Самое закрытое из устройств типа Coral Edge TPU вышло в марте 2019, и именно с ним мне удалось поиграться. Это один из прототипов Coral, хотя по документам у него абсолютно такая же производительность, как у выпущенных плат с mPCIe. С USB 3.0 проблем не возникает, но если вам не хочется возиться с внешним донглом, можно поставить в разъём платы M.2 Accelerator A+E вместо WiFi-карты:

Слева: моя WiFi-карта в ноутбуке. Справа: ускоритель M.2 A+E, вставляющийся в такой же слот.

И что же находится внутри краевого TPU-ускорителя размером 30 мм х 65 мм и стоимостью в $75? Может ли там быть mPCIe-ускоритель размером 30 мм х 22 мм за $35 с USB-адаптером и радиатором? Можем ли мы аккуратно вскрыть коробку из пластика и металла?

Передняя и задняя часть Edge TPU

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

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

Ответ на поставленный выше вопрос: нет, не буквально, но платы весьма схожи:

Рабочая сторона Coral Edge TPU. Справа – ускоритель M2 A+E и отдельный TPU-чип на одноцентовой монете

Чипы, касающиеся радиатора – это, вероятно, TPU (покрупнее) и память. Собрав Edge TPU обратно, я провёл собственные тесты на скорость, хотя тесты от Google говорят о 10- и даже 20-кратном ускорении работы по сравнению с CPU.

Пример с сайта Coral для распознавания объектов можно запускать и на CPU, и на TPU. Я подсоединяю это устройство к своему Acer Chromebook 11 с Galium OS 2.1, которая достаточно похожа на Debian для того, чтобы можно было без проблем установить tf-lite и программу для работы с TPU. Модели MobileNetV2 SSDLite для распознавания объектов занимают менее 7 МБ и работают с изображениями размером 300х300. На TPU одно изображение обрабатывается за 20 мс (+10 мс на однократное копирование изображения). На моём ноутбуке с процессором Intel Celeron 2.16GHz CPU 2014 года изображение обрабатывается за 1500 мс. На другом ПК, с процессором Intel Xeon 2.5 GHz CPU (который в 15 раз быстрее моего по данным cpubenchmark.net), изображение обрабатывается за 130 мс (при потреблении в 210 Вт!). По данным cpubenchmark.net, самый быстрый CPU из доступных будет в 3 раза быстрее этого: и это всё равно будет медленнее краевого TPU, который, к тому же, гораздо дешевле и потребляет гораздо меньше!

Более интересный проект – распознавание объектов в видеопотоке с камеры моего ноутбука. На процессоре Celeron выбранный пример успевает обработать меньше одного кадра в секунду. На TPU он обрабатывает 20 кадров в секунду – то есть, работает в реальном времени. И лишь четверть времени уходит на общение с TPU, остальное – это изменение размера картинок и вывод итоговых данных, которыми занимается CPU. Хотя я установил программу для работы TPU на максимальной частоте, я и близко не подошёл к пределам его возможностей, поэтому радиатор TPU не нагрелся до какой-либо некомфортной температуры.

Приведённые примеры было достаточно легко изменить и подстроить. Для начала работы с устройством нужно всего лишь установить рабочее окружение TPU через apt-get, потом установить tf-lite через pip – и всё готово для работы с примерами. Модели в этих примерах уже предварительно откомпилированы. Для более правдоподобного теста скорости нужно было бы обучить модель с нуля, запустить её на GPU и CPU, преобразовать в tf-lite, и выполнить его на CPU и TPU.

В целом, меня обрадовало это устройство. Мне кажется удивительным, что сеть, распознающая объекты на изображениях размера 300х300, работает на нетбуке 2014 года в реальном времени, с апгрейдом стоимостью всего в $75 (или даже $35, если взять плату M.2). Конечно, краевой TPU не заменит GPU для обучения и даже для оценки моделей. Однако я вижу большой потенциал у платы, в которой один TOPs обходится в $10 и потребляет 0,5 Вт. Это, по меньшей мере, в 10 раз дешевле, чем работать с GPU! И гораздо менее затратно по энергии. Это прекрасный вариант для демонстрации работы где-то в полях. Если интересующая вас задача указана в приведённой на сайте продукта табличке, то этот чип окажется полезным в случае, если приоритетом для вас будут низкое энергопотребление и небольшой размер системы.