Примечание.
Эта статья также доступна здесь. (на японском языке)
https://zenn.dev/wok/articles/0006_convert-from-pytorch-to-tfjs
Во-первых
Я хочу легко запускать модели машинного обучения в браузере. Поэтому я работаю над библиотекой для запуска различных моделей в tensorflowjs (см. этот репозиторий). Однако, как вы можете видеть на графике ниже, PyTorch в последнее время использовался в 75% докладов на крупных конференциях, что делает невозможным простой запуск сложных и интересных моделей с помощью tensorflowjs.

В этой ситуации PINTO, авторитет в области квантования моделей, опубликовал метод преобразования Pytorch в модели Tensorflow в Адвент-календаре этого года. Спасибо за очень замечательную статью.
Когда я попытался воспроизвести преобразование модели с помощью опубликованного метода, в моей среде были некоторые застревания, поэтому я решил сохранить эту статью в качестве рабочей заметки. В общем, надеюсь, вы ознакомитесь с оригинальной статьей ПИНТО. Если вы застряли, я был бы признателен, если бы вы могли взглянуть на это просто, чтобы увидеть, может ли это быть полезным.
Подход
Как я упоминал выше, некоторые вещи не работали из-за моей конкретной среды. Я пробовал много проб и ошибок, но это не сработало. Итак, мне наконец удалось воспроизвести его с помощью докера. Далее я буду исходить из того, что буду работать с докером.
Среда сборки
Сначала запустите контейнер докеров. Я использовал образ 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, поэтому установите его, обратившись к следующему сайту.
Этот сайт попросит вас загрузить пакет инструментов OpenVINO со следующего сайта
Вроде можно установить с помощью 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. Я не знаю, как его переключить…)

В следующей статье я также попробую портретный рисунок U²-Net. Это будет выглядеть так.


Ссылка
Я использовал изображения со следующих страниц.