Преобразование модели PyTorch в tensorflowjs

Архив Bezdepbonusu727.06.20225 мин чтения

Примечание.
Эта статья также доступна здесь. (на японском языке)
https://zenn.dev/wok/articles/0006_convert-from-pytorch-to-tfjs

Во-первых

Я хочу легко запускать модели машинного обучения в браузере. Поэтому я работаю над библиотекой для запуска различных моделей в tensorflowjs (см. этот репозиторий). Однако, как вы можете видеть на графике ниже, PyTorch в последнее время использовался в 75% докладов на крупных конференциях, что делает невозможным простой запуск сложных и интересных моделей с помощью tensorflowjs.

В этой ситуации PINTO, авторитет в области квантования моделей, опубликовал метод преобразования Pytorch в модели Tensorflow в Адвент-календаре этого года. Спасибо за очень замечательную статью.

Когда я попытался воспроизвести преобразование модели с помощью опубликованного метода, в моей среде были некоторые застревания, поэтому я решил сохранить эту статью в качестве рабочей заметки. В общем, надеюсь, вы ознакомитесь с оригинальной статьей ПИНТО. Если вы застряли, я был бы признателен, если бы вы могли взглянуть на это просто, чтобы увидеть, может ли это быть полезным.

[Русский] Моментальное преобразование моделей PyTorch, ONNX, Caffe и OpenVINO (NCHW) в Tensorflow / TensorflowLite (NHWC)

Подход

Как я упоминал выше, некоторые вещи не работали из-за моей конкретной среды. Я пробовал много проб и ошибок, но это не сработало. Итак, мне наконец удалось воспроизвести его с помощью докера. Далее я буду исходить из того, что буду работать с докером.

Среда сборки

Сначала запустите контейнер докеров. Я использовал образ Ubuntu с версией Cuda 10.1.

$ docker run --gpus all  -v /home/whoami/docker_share:/work --name converter -ti nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

Затем установите соответствующие пакеты. Я добавил emacs и mlocate по своему усмотрению, но если они вам не нужны, пропустите их.

$ apt install -y python3.8 python3-pip emacs mlocate pciutils cpio git sudo curl
$ python3 -m pip install pip --upgrade
$ pip3 install tensorflow==2.3.1 --upgrade
$ pip3 install openvino2tensorflow --upgrade
$ pip3 install torch==1.7.0+cu101 torchvision==0.8.1+cu101 torchaudio==0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
$ pip3 install onnxruntime onnx-simplifier openvino-python --upgrade
$ pip3 install networkx defusedxml test-generator==0.1.1 tensorflow_datasets tensorflowjs

Набор инструментов OpenVino, похоже, не включен, если вы просто используете pip, поэтому установите его, обратившись к следующему сайту.

Установите Intel® Distribution of OpenVINO™ toolkit для Linux* - OpenVINO™ Toolkit
ПРИМЕЧАНИЯ. Эти шаги применимы к Ubuntu*, CentOS* и Yocto*. Если вы используете Intel® Distribution of OpenVINO™ toolkit на...docs.openvinotoolkit.org

Этот сайт попросит вас загрузить пакет инструментов OpenVINO со следующего сайта

Загрузить дистрибутив Intel® OpenVINO™ Toolkit
Загрузить версию дистрибутива Intel® набора инструментов OpenVINO™ для Linux, Windows или macOS.software.intel .com

Вроде можно установить с помощью pip, но когда я устанавливал с помощью pip, некоторые команды не были найдены в моем окружении и я не мог продолжить работу.

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

Выберите 2021.1 в качестве версии и загрузите полный пакет. После загрузки поместите файл в место, которое видно из контейнера Docker. В данном случае я положил его в /work/l_openvino_toolkit_p_2021.1.110.tgz.

Теперь давайте установим его.
Когда начнется установка, вам будет предложен выбор. Будут различные предупреждения, но я пропущу предварительные условия.

$ cp /work/l_openvino_toolkit_p_2021.1.110.tgz ./
$ tar xvfz l_openvino_toolkit_p_2021.1.110.tgz
$ cd l_openvino_toolkit_p_2021.1.110
$ ./install.sh
$ cd -

Следующим шагом будет установка зависимых модулей.

$ cd /opt/intel/openvino_2021/install_dependencies
$ ./install_openvino_dependencies.sh
$ source /opt/intel/openvino_2021/bin/setupvars.sh
$ cd -

Вот и все!

Преобразование

Теперь давайте посмотрим, сможем ли мы правильно ее преобразовать.
Давайте попробуем преобразовать модель семантической сегментации U²-Net, скопировав статью ПИНТО.

Клонируйте репозиторий U²-Net и загрузите модель.

$ git clone https://github.com/NathanUA/U-2-Net.git
$ cd U-2-Net/
$ mkdir ./saved_models/u2netp/
$ curl -sc /tmp/cookie "https://drive.google.com/uc?export=download&id=1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy" > /dev/null
$ CODE="$(awk '/_warning_/ {print $NF}' /tmp/cookie)"
$ curl -Lb /tmp/cookie "https://drive.google.com/uc?export=download&confirm=${CODE}&id=1rbSTGKAE-MTxBYHd-51l2hMOQPT_7EPy" -o saved_models/u2netp/u2netp.pth

Начните процесс преобразования ниже.

$ export PYTHONPATH=/U-2-Net
$ SIZE=512
$ python3 /opt/intel/openvino_2021/deployment_tools/tools/model_downloader/pytorch_to_onnx.py  \
 --import-module model.u2net   \
 --model-name U2NETP   \
 --input-shape 1,3,${SIZE},${SIZE}   \
 --weights saved_models/u2netp/u2netp.pth   \
 --output-file u2netp_${SIZE}x${SIZE}.onnx   --input-names "x"   \
 --output-names "a/F.sigmoid(d0)"
<snip...>
ONNX check passed successfully.
$ python3 -m onnxsim u2netp_${SIZE}x${SIZE}.onnx u2netp_${SIZE}x${SIZE}_opt.onnx
<snip...>
Checking 2/3...
Ok!
$ python3 /opt/intel/openvino_2021/deployment_tools/model_optimizer/mo.py \
  --input_model u2netp_${SIZE}x${SIZE}_opt.onnx \
  --input_shape [1,3,${SIZE},${SIZE}] \
  --output_dir openvino/${SIZE}x${SIZE}/FP32 \
  --data_type FP32
<snip...>
[ SUCCESS ] Generated IR version 10 model.
[ SUCCESS ] XML file: /U-2-Net/openvino/512x512/FP32/u2netp_512x512.xml
[ SUCCESS ] BIN file: /U-2-Net/openvino/512x512/FP32/u2netp_512x512.bin
[ SUCCESS ] Total execution time: 35.76 seconds.
[ SUCCESS ] Memory consumed: 809 MB.
$ openvino2tensorflow \
    --model_path openvino/${SIZE}x${SIZE}/FP32/u2netp_${SIZE}x${SIZE}_opt.xml \
    --model_output_path saved_model_${SIZE}x${SIZE} \
    --output_saved_model True
<snip...>
All the conversion process is finished! =============================================
$ tensorflowjs_converter \
  --input_format=tf_saved_model \
  --output_format=tfjs_graph_model \
  --signature_name=serving_default \
  --saved_model_tags=serve \
  saved_model_${SIZE}x${SIZE} \
  tfjs_model_u2netp_${SIZE}x${SIZE}

Теперь у нас может быть модель Tensorflowjs!

Используй это

Давайте используем эту модель, чтобы фактически сегментировать изображение. Результат показан ниже. Мы видим, что сегментация работает. Он немного размыт в некоторых областях, но я думаю, что это можно улучшить, правильно установив пороговое значение.
Я слышал, что многие люди говорят, что он довольно легкий, но в моей среде (процессор) заняло 2-3 минуты. Кроме того, с графическим процессором, даже с GTX1660 с 6 ГБ ОЗУ, я не мог запустить его, потому что он переполнял память. (У меня также есть 2080Ti на той же машине, но Chrome активирует только GTX1660. Я не знаю, как его переключить...)