enhance armnn conversion

machine-learning/export/ann/.gitignore

@@ -0,0 +1,2 @@
+armnn*
+output/

machine-learning/export/ann/Dockerfile

@@ -0,0 +1,28 @@
+FROM mambaorg/micromamba:bookworm-slim@sha256:333f7598ff2c2400fb10bfe057709c68b7daab5d847143af85abcf224a07271a as builder
+
+ENV TRANSFORMERS_CACHE=/cache \
+  PYTHONDONTWRITEBYTECODE=1 \
+  PYTHONUNBUFFERED=1 \
+  PATH="/opt/venv/bin:$PATH"
+
+WORKDIR /export/ann
+
+USER root
+RUN apt-get update && apt-get install -y --no-install-recommends \
+  build-essential \
+  curl \
+  git
+
+USER $MAMBA_USER
+COPY --chown=$MAMBA_USER:$MAMBA_USER env.yaml ./
+RUN micromamba install -y -f env.yaml
+COPY --chown=$MAMBA_USER:$MAMBA_USER *.sh *.cpp ./
+
+ENV ARMNN_PATH=/export/ann/armnn
+RUN ./download-armnn.sh && \
+    ./build-converter.sh && \
+    ./build.sh
+COPY --chown=$MAMBA_USER:$MAMBA_USER run.py ./
+
+ENTRYPOINT ["/usr/local/bin/_entrypoint.sh"]
+CMD ["python", "run.py"]

machine-learning/export/ann/ann.cpp

@@ -0,0 +1,281 @@
+#include <fstream>
+#include <mutex>
+#include <atomic>
+
+#include "armnn/IRuntime.hpp"
+#include "armnn/INetwork.hpp"
+#include "armnn/Types.hpp"
+#include "armnnDeserializer/IDeserializer.hpp"
+#include "armnnTfLiteParser/ITfLiteParser.hpp"
+#include "armnnOnnxParser/IOnnxParser.hpp"
+
+using namespace armnn;
+
+struct IOInfos
+{
+    std::vector<BindingPointInfo> inputInfos;
+    std::vector<BindingPointInfo> outputInfos;
+};
+
+// from https://rigtorp.se/spinlock/
+struct SpinLock
+{
+    std::atomic<bool> lock_ = {false};
+
+    void lock()
+    {
+        for (;;)
+        {
+            if (!lock_.exchange(true, std::memory_order_acquire))
+            {
+                break;
+            }
+            while (lock_.load(std::memory_order_relaxed))
+                ;
+        }
+    }
+
+    void unlock() { lock_.store(false, std::memory_order_release); }
+};
+
+class Ann
+{
+
+public:
+    int load(const char *modelPath,
+             bool fastMath,
+             bool fp16,
+             bool saveCachedNetwork,
+             const char *cachedNetworkPath)
+    {
+        INetworkPtr network = loadModel(modelPath);
+        IOptimizedNetworkPtr optNet = OptimizeNetwork(network.get(), fastMath, fp16, saveCachedNetwork, cachedNetworkPath);
+        const IOInfos infos = getIOInfos(optNet.get());
+        NetworkId netId;
+        mutex.lock();
+        Status status = runtime->LoadNetwork(netId, std::move(optNet));
+        mutex.unlock();
+        if (status != Status::Success)
+        {
+            return -1;
+        }
+        spinLock.lock();
+        ioInfos[netId] = infos;
+        mutexes.emplace(netId, std::make_unique<std::mutex>());
+        spinLock.unlock();
+        return netId;
+    }
+
+    void execute(NetworkId netId, const void **inputData, void **outputData)
+    {
+        spinLock.lock();
+        const IOInfos *infos = &ioInfos[netId];
+        auto m = mutexes[netId].get();
+        spinLock.unlock();
+        InputTensors inputTensors;
+        inputTensors.reserve(infos->inputInfos.size());
+        size_t i = 0;
+        for (const BindingPointInfo &info : infos->inputInfos)
+            inputTensors.emplace_back(info.first, ConstTensor(info.second, inputData[i++]));
+        OutputTensors outputTensors;
+        outputTensors.reserve(infos->outputInfos.size());
+        i = 0;
+        for (const BindingPointInfo &info : infos->outputInfos)
+            outputTensors.emplace_back(info.first, Tensor(info.second, outputData[i++]));
+        m->lock();
+        runtime->EnqueueWorkload(netId, inputTensors, outputTensors);
+        m->unlock();
+    }
+
+    void unload(NetworkId netId)
+    {
+        mutex.lock();
+        runtime->UnloadNetwork(netId);
+        mutex.unlock();
+    }
+
+    int tensors(NetworkId netId, bool isInput = false)
+    {
+        spinLock.lock();
+        const IOInfos *infos = &ioInfos[netId];
+        spinLock.unlock();
+        return (int)(isInput ? infos->inputInfos.size() : infos->outputInfos.size());
+    }
+
+    unsigned long shape(NetworkId netId, bool isInput = false, int index = 0)
+    {
+        spinLock.lock();
+        const IOInfos *infos = &ioInfos[netId];
+        spinLock.unlock();
+        const TensorShape shape = (isInput ? infos->inputInfos : infos->outputInfos)[index].second.GetShape();
+        unsigned long s = 0;
+        for (unsigned int d = 0; d < shape.GetNumDimensions(); d++)
+            s |= ((unsigned long)shape[d]) << (d * 16); // stores up to 4 16-bit values in a 64-bit value
+        return s;
+    }
+
+    Ann(int tuningLevel, const char *tuningFile)
+    {
+        IRuntime::CreationOptions runtimeOptions;
+        BackendOptions backendOptions{"GpuAcc",
+                                      {
+                                          {"TuningLevel", tuningLevel},
+                                          {"MemoryOptimizerStrategy", "ConstantMemoryStrategy"}, // SingleAxisPriorityList or ConstantMemoryStrategy
+                                      }};
+        if (tuningFile)
+            backendOptions.AddOption({"TuningFile", tuningFile});
+        runtimeOptions.m_BackendOptions.emplace_back(backendOptions);
+        runtime = IRuntime::CreateRaw(runtimeOptions);
+    };
+    ~Ann()
+    {
+        IRuntime::Destroy(runtime);
+    };
+
+private:
+    INetworkPtr loadModel(const char *modelPath)
+    {
+        const auto path = std::string(modelPath);
+        if (path.rfind(".tflite") == path.length() - 7) // endsWith()
+        {
+            auto parser = armnnTfLiteParser::ITfLiteParser::CreateRaw();
+            return parser->CreateNetworkFromBinaryFile(modelPath);
+        }
+        else if (path.rfind(".onnx") == path.length() - 5) // endsWith()
+        {
+            auto parser = armnnOnnxParser::IOnnxParser::CreateRaw();
+            return parser->CreateNetworkFromBinaryFile(modelPath);
+        }
+        else
+        {
+            std::ifstream ifs(path, std::ifstream::in | std::ifstream::binary);
+            auto parser = armnnDeserializer::IDeserializer::CreateRaw();
+            return parser->CreateNetworkFromBinary(ifs);
+        }
+    }
+
+    static BindingPointInfo getInputTensorInfo(LayerBindingId inputBindingId, TensorInfo info)
+    {
+        const auto newInfo = TensorInfo{info.GetShape(), info.GetDataType(),
+                                        info.GetQuantizationScale(),
+                                        info.GetQuantizationOffset(),
+                                        true};
+        return {inputBindingId, newInfo};
+    }
+
+    IOptimizedNetworkPtr OptimizeNetwork(INetwork *network, bool fastMath, bool fp16, bool saveCachedNetwork, const char *cachedNetworkPath)
+    {
+        const bool allowExpandedDims = false;
+        const ShapeInferenceMethod shapeInferenceMethod = ShapeInferenceMethod::ValidateOnly;
+
+        OptimizerOptionsOpaque options;
+        options.SetReduceFp32ToFp16(fp16);
+        options.SetShapeInferenceMethod(shapeInferenceMethod);
+        options.SetAllowExpandedDims(allowExpandedDims);
+
+        BackendOptions gpuAcc("GpuAcc", {{"FastMathEnabled", fastMath}});
+        if (cachedNetworkPath)
+        {
+            gpuAcc.AddOption({"SaveCachedNetwork", saveCachedNetwork});
+            gpuAcc.AddOption({"CachedNetworkFilePath", cachedNetworkPath});
+        }
+        options.AddModelOption(gpuAcc);
+
+        // No point in using ARMNN for CPU, use ONNX (quantized) instead.
+        // BackendOptions cpuAcc("CpuAcc",
+        //                       {
+        //                           {"FastMathEnabled", fastMath},
+        //                           {"NumberOfThreads", 0},
+        //                       });
+        // options.AddModelOption(cpuAcc);
+
+        BackendOptions allowExDimOpt("AllowExpandedDims",
+                                     {{"AllowExpandedDims", allowExpandedDims}});
+        options.AddModelOption(allowExDimOpt);
+        BackendOptions shapeInferOpt("ShapeInferenceMethod",
+                                     {{"InferAndValidate", shapeInferenceMethod == ShapeInferenceMethod::InferAndValidate}});
+        options.AddModelOption(shapeInferOpt);
+
+        std::vector<BackendId> backends = {
+            BackendId("GpuAcc"),
+            // BackendId("CpuAcc"),
+            // BackendId("CpuRef"),
+        };
+        return Optimize(*network, backends, runtime->GetDeviceSpec(), options);
+    }
+
+    IOInfos getIOInfos(IOptimizedNetwork *optNet)
+    {
+        struct InfoStrategy : IStrategy
+        {
+            void ExecuteStrategy(const IConnectableLayer *layer,
+                                 const BaseDescriptor &descriptor,
+                                 const std::vector<ConstTensor> &constants,
+                                 const char *name,
+                                 const LayerBindingId id = 0) override
+            {
+                IgnoreUnused(descriptor, constants, id);
+                const LayerType lt = layer->GetType();
+                if (lt == LayerType::Input)
+                    ioInfos.inputInfos.push_back(getInputTensorInfo(id, layer->GetOutputSlot(0).GetTensorInfo()));
+                else if (lt == LayerType::Output)
+                    ioInfos.outputInfos.push_back({id, layer->GetInputSlot(0).GetTensorInfo()});
+            }
+            IOInfos ioInfos;
+        };
+
+        InfoStrategy infoStrategy;
+        optNet->ExecuteStrategy(infoStrategy);
+        return infoStrategy.ioInfos;
+    }
+
+    IRuntime *runtime;
+    std::map<NetworkId, IOInfos> ioInfos;
+    std::map<NetworkId, std::unique_ptr<std::mutex>> mutexes; // mutex per network to not execute the same the same network concurrently
+    std::mutex mutex; // global mutex for load/unload calls to the runtime
+    SpinLock spinLock; // fast spin lock to guard access to the ioInfos and mutexes maps
+};
+
+extern "C" void *init(int logLevel, int tuningLevel, const char *tuningFile)
+{
+    LogSeverity level = static_cast<LogSeverity>(logLevel);
+    ConfigureLogging(true, true, level);
+
+    Ann *ann = new Ann(tuningLevel, tuningFile);
+    return ann;
+}
+
+extern "C" void destroy(void *ann)
+{
+    delete ((Ann *)ann);
+}
+
+extern "C" int load(void *ann,
+                    const char *path,
+                    bool fastMath,
+                    bool fp16,
+                    bool saveCachedNetwork,
+                    const char *cachedNetworkPath)
+{
+    return ((Ann *)ann)->load(path, fastMath, fp16, saveCachedNetwork, cachedNetworkPath);
+}
+
+extern "C" void unload(void *ann, NetworkId netId)
+{
+    ((Ann *)ann)->unload(netId);
+}
+
+extern "C" void execute(void *ann, NetworkId netId, const void **inputData, void **outputData)
+{
+    ((Ann *)ann)->execute(netId, inputData, outputData);
+}
+
+extern "C" unsigned long shape(void *ann, NetworkId netId, bool isInput, int index)
+{
+    return ((Ann *)ann)->shape(netId, isInput, index);
+}
+
+extern "C" int tensors(void *ann, NetworkId netId, bool isInput)
+{
+    return ((Ann *)ann)->tensors(netId, isInput);
+}

machine-learning/export/ann/build-converter.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env sh
+
+cd armnn-23.11/ || exit
+g++ -o ../armnnconverter -fPIC -O1 -DARMNN_ONNX_PARSER -DARMNN_SERIALIZER -DARMNN_TF_LITE_PARSER -fuse-ld=gold -std=c++17 -Iinclude -Isrc/armnnUtils -Ithird-party -larmnn -larmnnDeserializer -larmnnTfLiteParser -larmnnOnnxParser -larmnnSerializer -L../armnn src/armnnConverter/ArmnnConverter.cpp

machine-learning/export/ann/build.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env sh
+
+g++ -shared -O3 -fPIC -o libann.so -fuse-ld=gold -std=c++17 -I"$ARMNN_PATH"/include -larmnn -larmnnDeserializer -larmnnTfLiteParser -larmnnOnnxParser -L"$ARMNN_PATH" ann.cpp

machine-learning/export/ann/download-armnn.sh

@@ -0,0 +1,8 @@
+#!/bin/sh
+
+# binaries
+mkdir armnn
+curl -SL "https://github.com/ARM-software/armnn/releases/download/v23.11/ArmNN-linux-x86_64.tar.gz" | tar -zx -C armnn
+
+# source to build ArmnnConverter
+curl -SL "https://github.com/ARM-software/armnn/archive/refs/tags/v23.11.tar.gz" | tar -zx

machine-learning/export/ann/env.yaml

@@ -0,0 +1,201 @@
+name: annexport
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+dependencies:
+  - _libgcc_mutex=0.1=conda_forge
+  - _openmp_mutex=4.5=2_kmp_llvm
+  - aiohttp=3.9.1=py310h2372a71_0
+  - aiosignal=1.3.1=pyhd8ed1ab_0
+  - arpack=3.8.0=nompi_h0baa96a_101
+  - async-timeout=4.0.3=pyhd8ed1ab_0
+  - attrs=23.1.0=pyh71513ae_1
+  - aws-c-auth=0.7.3=h28f7589_1
+  - aws-c-cal=0.6.1=hc309b26_1
+  - aws-c-common=0.9.0=hd590300_0
+  - aws-c-compression=0.2.17=h4d4d85c_2
+  - aws-c-event-stream=0.3.1=h2e3709c_4
+  - aws-c-http=0.7.11=h00aa349_4
+  - aws-c-io=0.13.32=he9a53bd_1
+  - aws-c-mqtt=0.9.3=hb447be9_1
+  - aws-c-s3=0.3.14=hf3aad02_1
+  - aws-c-sdkutils=0.1.12=h4d4d85c_1
+  - aws-checksums=0.1.17=h4d4d85c_1
+  - aws-crt-cpp=0.21.0=hb942446_5
+  - aws-sdk-cpp=1.10.57=h85b1a90_19
+  - blas=2.120=openblas
+  - blas-devel=3.9.0=20_linux64_openblas
+  - brotli-python=1.0.9=py310hd8f1fbe_9
+  - bzip2=1.0.8=hd590300_5
+  - c-ares=1.23.0=hd590300_0
+  - ca-certificates=2023.11.17=hbcca054_0
+  - certifi=2023.11.17=pyhd8ed1ab_0
+  - charset-normalizer=3.3.2=pyhd8ed1ab_0
+  - click=8.1.7=unix_pyh707e725_0
+  - colorama=0.4.6=pyhd8ed1ab_0
+  - coloredlogs=15.0.1=pyhd8ed1ab_3
+  - cuda-cudart=11.7.99=0
+  - cuda-cupti=11.7.101=0
+  - cuda-libraries=11.7.1=0
+  - cuda-nvrtc=11.7.99=0
+  - cuda-nvtx=11.7.91=0
+  - cuda-runtime=11.7.1=0
+  - dataclasses=0.8=pyhc8e2a94_3
+  - datasets=2.14.7=pyhd8ed1ab_0
+  - dill=0.3.7=pyhd8ed1ab_0
+  - filelock=3.13.1=pyhd8ed1ab_0
+  - flatbuffers=23.5.26=h59595ed_1
+  - freetype=2.12.1=h267a509_2
+  - frozenlist=1.4.0=py310h2372a71_1
+  - fsspec=2023.10.0=pyhca7485f_0
+  - ftfy=6.1.3=pyhd8ed1ab_0
+  - gflags=2.2.2=he1b5a44_1004
+  - glog=0.6.0=h6f12383_0
+  - glpk=5.0=h445213a_0
+  - gmp=6.3.0=h59595ed_0
+  - gmpy2=2.1.2=py310h3ec546c_1
+  - huggingface_hub=0.17.3=pyhd8ed1ab_0
+  - humanfriendly=10.0=pyhd8ed1ab_6
+  - icu=73.2=h59595ed_0
+  - idna=3.6=pyhd8ed1ab_0
+  - importlib-metadata=7.0.0=pyha770c72_0
+  - importlib_metadata=7.0.0=hd8ed1ab_0
+  - joblib=1.3.2=pyhd8ed1ab_0
+  - keyutils=1.6.1=h166bdaf_0
+  - krb5=1.21.2=h659d440_0
+  - lcms2=2.15=h7f713cb_2
+  - ld_impl_linux-64=2.40=h41732ed_0
+  - lerc=4.0.0=h27087fc_0
+  - libabseil=20230125.3=cxx17_h59595ed_0
+  - libarrow=12.0.1=hb87d912_8_cpu
+  - libblas=3.9.0=20_linux64_openblas
+  - libbrotlicommon=1.0.9=h166bdaf_9
+  - libbrotlidec=1.0.9=h166bdaf_9
+  - libbrotlienc=1.0.9=h166bdaf_9
+  - libcblas=3.9.0=20_linux64_openblas
+  - libcrc32c=1.1.2=h9c3ff4c_0
+  - libcublas=11.10.3.66=0
+  - libcufft=10.7.2.124=h4fbf590_0
+  - libcufile=1.8.1.2=0
+  - libcurand=10.3.4.101=0
+  - libcurl=8.5.0=hca28451_0
+  - libcusolver=11.4.0.1=0
+  - libcusparse=11.7.4.91=0
+  - libdeflate=1.19=hd590300_0
+  - libedit=3.1.20191231=he28a2e2_2
+  - libev=4.33=hd590300_2
+  - libevent=2.1.12=hf998b51_1
+  - libffi=3.4.2=h7f98852_5
+  - libgcc-ng=13.2.0=h807b86a_3
+  - libgfortran-ng=13.2.0=h69a702a_3
+  - libgfortran5=13.2.0=ha4646dd_3
+  - libgoogle-cloud=2.12.0=hac9eb74_1
+  - libgrpc=1.54.3=hb20ce57_0
+  - libhwloc=2.9.3=default_h554bfaf_1009
+  - libiconv=1.17=hd590300_1
+  - libjpeg-turbo=2.1.5.1=hd590300_1
+  - liblapack=3.9.0=20_linux64_openblas
+  - liblapacke=3.9.0=20_linux64_openblas
+  - libnghttp2=1.58.0=h47da74e_1
+  - libnpp=11.7.4.75=0
+  - libnsl=2.0.1=hd590300_0
+  - libnuma=2.0.16=h0b41bf4_1
+  - libnvjpeg=11.8.0.2=0
+  - libopenblas=0.3.25=pthreads_h413a1c8_0
+  - libpng=1.6.39=h753d276_0
+  - libprotobuf=3.21.12=hfc55251_2
+  - libsentencepiece=0.1.99=h180e1df_0
+  - libsqlite=3.44.2=h2797004_0
+  - libssh2=1.11.0=h0841786_0
+  - libstdcxx-ng=13.2.0=h7e041cc_3
+  - libthrift=0.18.1=h8fd135c_2
+  - libtiff=4.6.0=h29866fb_1
+  - libutf8proc=2.8.0=h166bdaf_0
+  - libuuid=2.38.1=h0b41bf4_0
+  - libwebp-base=1.3.2=hd590300_0
+  - libxcb=1.15=h0b41bf4_0
+  - libxml2=2.11.6=h232c23b_0
+  - libzlib=1.2.13=hd590300_5
+  - llvm-openmp=17.0.6=h4dfa4b3_0
+  - lz4-c=1.9.4=hcb278e6_0
+  - mkl=2022.2.1=h84fe81f_16997
+  - mkl-devel=2022.2.1=ha770c72_16998
+  - mkl-include=2022.2.1=h84fe81f_16997
+  - mpc=1.3.1=hfe3b2da_0
+  - mpfr=4.2.1=h9458935_0
+  - mpmath=1.3.0=pyhd8ed1ab_0
+  - multidict=6.0.4=py310h2372a71_1
+  - multiprocess=0.70.15=py310h2372a71_1
+  - ncurses=6.4=h59595ed_2
+  - numpy=1.26.2=py310hb13e2d6_0
+  - onnx=1.14.0=py310ha3deec4_1
+  - onnx2torch=1.5.13=pyhd8ed1ab_0
+  - onnxruntime=1.16.3=py310hd4b7fbc_1_cpu
+  - open-clip-torch=2.23.0=pyhd8ed1ab_1
+  - openblas=0.3.25=pthreads_h7a3da1a_0
+  - openjpeg=2.5.0=h488ebb8_3
+  - openssl=3.2.0=hd590300_1
+  - orc=1.9.0=h2f23424_1
+  - packaging=23.2=pyhd8ed1ab_0
+  - pandas=2.1.4=py310hcc13569_0
+  - pillow=10.0.1=py310h29da1c1_1
+  - pip=23.3.1=pyhd8ed1ab_0
+  - protobuf=4.21.12=py310heca2aa9_0
+  - pthread-stubs=0.4=h36c2ea0_1001
+  - pyarrow=12.0.1=py310h0576679_8_cpu
+  - pyarrow-hotfix=0.6=pyhd8ed1ab_0
+  - pysocks=1.7.1=pyha2e5f31_6
+  - python=3.10.13=hd12c33a_0_cpython
+  - python-dateutil=2.8.2=pyhd8ed1ab_0
+  - python-flatbuffers=23.5.26=pyhd8ed1ab_0
+  - python-tzdata=2023.3=pyhd8ed1ab_0
+  - python-xxhash=3.4.1=py310h2372a71_0
+  - python_abi=3.10=4_cp310
+  - pytorch=1.13.1=cpu_py310hd11e9c7_1
+  - pytorch-cuda=11.7=h778d358_5
+  - pytorch-mutex=1.0=cuda
+  - pytz=2023.3.post1=pyhd8ed1ab_0
+  - pyyaml=6.0.1=py310h2372a71_1
+  - rdma-core=28.9=h59595ed_1
+  - re2=2023.03.02=h8c504da_0
+  - readline=8.2=h8228510_1
+  - regex=2023.10.3=py310h2372a71_0
+  - requests=2.31.0=pyhd8ed1ab_0
+  - s2n=1.3.49=h06160fa_0
+  - sacremoses=0.0.53=pyhd8ed1ab_0
+  - safetensors=0.3.3=py310hcb5633a_1
+  - sentencepiece=0.1.99=hff52083_0
+  - sentencepiece-python=0.1.99=py310hebdb9f0_0
+  - sentencepiece-spm=0.1.99=h180e1df_0
+  - setuptools=68.2.2=pyhd8ed1ab_0
+  - six=1.16.0=pyh6c4a22f_0
+  - sleef=3.5.1=h9b69904_2
+  - snappy=1.1.10=h9fff704_0
+  - sympy=1.12=pypyh9d50eac_103
+  - tbb=2021.11.0=h00ab1b0_0
+  - texttable=1.7.0=pyhd8ed1ab_0
+  - timm=0.9.12=pyhd8ed1ab_0
+  - tk=8.6.13=noxft_h4845f30_101
+  - tokenizers=0.14.1=py310h320607d_2
+  - torchvision=0.14.1=cpu_py310hd3d2ac3_1
+  - tqdm=4.66.1=pyhd8ed1ab_0
+  - transformers=4.35.2=pyhd8ed1ab_0
+  - typing-extensions=4.9.0=hd8ed1ab_0
+  - typing_extensions=4.9.0=pyha770c72_0
+  - tzdata=2023c=h71feb2d_0
+  - ucx=1.14.1=h64cca9d_5
+  - urllib3=2.1.0=pyhd8ed1ab_0
+  - wcwidth=0.2.12=pyhd8ed1ab_0
+  - wheel=0.42.0=pyhd8ed1ab_0
+  - xorg-libxau=1.0.11=hd590300_0
+  - xorg-libxdmcp=1.1.3=h7f98852_0
+  - xxhash=0.8.2=hd590300_0
+  - xz=5.2.6=h166bdaf_0
+  - yaml=0.2.5=h7f98852_2
+  - yarl=1.9.3=py310h2372a71_0
+  - zipp=3.17.0=pyhd8ed1ab_0
+  - zlib=1.2.13=hd590300_5
+  - zstd=1.5.5=hfc55251_0
+  - pip:
+      - git+https://github.com/fyfrey/TinyNeuralNetwork.git

machine-learning/export/ann/run.py

@@ -0,0 +1,297 @@
+import os
+import platform
+import subprocess
+from tempfile import TemporaryDirectory
+from typing import Callable, ClassVar
+
+import onnx
+import torch
+from onnx2torch import convert
+from onnx2torch.node_converters.registry import add_converter
+from onnxruntime.tools.onnx_model_utils import fix_output_shapes, make_input_shape_fixed
+from tinynn.converter import TFLiteConverter
+from huggingface_hub import snapshot_download
+from onnx2torch.onnx_graph import OnnxGraph
+from onnx2torch.onnx_node import OnnxNode
+from onnx2torch.utils.common import OperationConverterResult, onnx_mapping_from_node
+from onnx.shape_inference import infer_shapes_path
+from huggingface_hub import login, upload_file
+
+# egregious hacks:
+# changed `Clip`'s min/max logic to skip empty strings
+# changed OnnxSqueezeDynamicAxes to use `sorted` instead of `torch.sort``
+# commented out shape inference in `fix_output_shapes``
+
+
+class ArgMax(torch.nn.Module):
+    def __init__(self, dim: int = -1, keepdim: bool = False):
+        super().__init__()
+        self.dim = dim
+        self.keepdim = keepdim
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return torch.argmax(input, dim=self.dim, keepdim=self.keepdim)
+
+
+class Erf(torch.nn.Module):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return torch.erf(input)
+
+
+@add_converter(operation_type="ArgMax", version=13)
+def _(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:
+    return OperationConverterResult(
+        torch_module=ArgMax(),
+        onnx_mapping=onnx_mapping_from_node(node=node),
+    )
+
+
+class ExportBase(torch.nn.Module):
+    task: ClassVar[str]
+
+    def __init__(
+        self,
+        name: str,
+        input_shape: tuple[int, ...],
+        pretrained: str | None = None,
+        optimization_level: int = 5,
+    ):
+        super().__init__()
+        self.name = name
+        self.optimize = optimization_level
+        self.nchw_transpose = False
+        self.input_shape = input_shape
+        self.pretrained = pretrained
+        self.dummy_param = torch.nn.Parameter(torch.empty(0))
+        self.model = self.load().eval()
+        for param in self.parameters():
+            param.requires_grad_(False)
+        self.eval()
+
+    def load(self) -> torch.nn.Module:
+        cache_dir = os.path.join(os.environ["CACHE_DIR"], self.model_name)
+        task_path = os.path.join(cache_dir, self.task)
+        model_path = os.path.join(task_path, "model.onnx")
+        if not os.path.isfile(model_path):
+            snapshot_download(self.repo_name, cache_dir=cache_dir, local_dir=cache_dir)
+        infer_shapes_path(model_path, check_type=True, strict_mode=True, data_prop=True)
+        onnx_model = onnx.load_model(model_path)
+        make_input_shape_fixed(onnx_model.graph, onnx_model.graph.input[0].name, self.input_shape)
+        fix_output_shapes(onnx_model)
+        # try:
+        # onnx.save(onnx_model, model_path)
+        # except:
+        # onnx.save(onnx_model, model_path, save_as_external_data=True, all_tensors_to_one_file=False)
+        # infer_shapes_path(model_path, check_type=True, strict_mode=True, data_prop=True)
+        # onnx_model = onnx.load_model(model_path)
+        # onnx_model = infer_shapes(onnx_model, check_type=True, strict_mode=True, data_prop=True)
+        return convert(onnx_model)
+
+    def forward(self, *inputs: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor]:
+        if self.precision == "fp16":
+            inputs = tuple(i.half() for i in inputs)
+
+        out = self._forward(*inputs)
+        if self.precision == "fp16":
+            if isinstance(out, tuple):
+                return tuple(o.float() for o in out)
+            return out.float()
+        return out
+
+    def _forward(self, *inputs: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor]:
+        return self.model(*inputs)
+
+    def to_armnn(self, output_path: str) -> None:
+        output_dir = os.path.dirname(output_path)
+        os.makedirs(output_dir, exist_ok=True)
+        self(*self.dummy_inputs)
+        print(f"Exporting {self.model_name} ({self.task}) with {self.precision} precision")
+        jit = torch.jit.trace(self, self.dummy_inputs).eval()
+        with TemporaryDirectory() as tmpdir:
+            tflite_model_path = os.path.join(tmpdir, "model.tflite")
+            converter = TFLiteConverter(
+                jit,
+                self.dummy_inputs,
+                tflite_model_path,
+                optimize=self.optimize,
+                nchw_transpose=self.nchw_transpose,
+            )
+            # segfaults on ARM, must run on x86_64 / AMD64
+            converter.convert()
+
+            subprocess.run(
+                [
+                    "./armnnconverter",
+                    "-f",
+                    "tflite-binary",
+                    "-m",
+                    tflite_model_path,
+                    "-i",
+                    "input_tensor",
+                    "-o",
+                    "output_tensor",
+                    "-p",
+                    output_path,
+                ],
+                capture_output=True,
+            )
+        print(f"Finished exporting {self.name} ({self.task}) with {self.precision} precision")
+
+    @property
+    def dummy_inputs(self) -> tuple[torch.FloatTensor]:
+        return (torch.rand(self.input_shape, device=self.device, dtype=self.dtype),)
+
+    @property
+    def model_name(self) -> str:
+        return f"{self.name}__{self.pretrained}" if self.pretrained else self.name
+
+    @property
+    def repo_name(self) -> str:
+        return f"immich-app/{self.model_name}"
+
+    @property
+    def device(self) -> torch.device:
+        return self.dummy_param.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.dummy_param.dtype
+
+    @property
+    def precision(self) -> str:
+        match self.dtype:
+            case torch.float32:
+                return "fp32"
+            case torch.float16:
+                return "fp16"
+            case _:
+                raise ValueError(f"Unsupported dtype {self.dtype}")
+
+
+class ArcFace(ExportBase):
+    task = "recognition"
+
+
+class RetinaFace(ExportBase):
+    task = "detection"
+
+
+class OpenClipVisual(ExportBase):
+    task = "visual"
+
+
+class OpenClipTextual(ExportBase):
+    task = "textual"
+
+    @property
+    def dummy_inputs(self) -> tuple[torch.LongTensor]:
+        return (torch.randint(0, 5000, self.input_shape, device=self.device, dtype=torch.int32),)
+
+
+class MClipTextual(ExportBase):
+    task = "textual"
+
+    @property
+    def dummy_inputs(self) -> tuple[torch.LongTensor]:
+        return (
+            torch.randint(0, 5000, self.input_shape, device=self.device, dtype=torch.int32),
+            torch.randint(0, 1, self.input_shape, device=self.device, dtype=torch.int32),
+        )
+
+
+def main() -> None:
+    if platform.machine() not in ("x86_64", "AMD64"):
+        raise RuntimeError(f"Can only run on x86_64 / AMD64, not {platform.machine()}")
+    login(token=os.environ["HF_AUTH_TOKEN"])
+    os.environ["LD_LIBRARY_PATH"] = "armnn"
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    failed: list[Callable[[], ExportBase]] = [
+        lambda: OpenClipVisual("ViT-H-14-378-quickgelu", (1, 3, 378, 378), pretrained="dfn5b"), # flatbuffers: cannot grow buffer beyond 2 gigabytes (will probably work with fp16)
+        lambda: OpenClipVisual("ViT-H-14-quickgelu", (1, 3, 224, 224), pretrained="dfn5b"), # flatbuffers: cannot grow buffer beyond 2 gigabytes (will probably work with fp16)
+        lambda: OpenClipTextual("nllb-clip-base-siglip", (1, 77), pretrained="v1"), # ERROR (tinynn.converter.base) Unsupported ops: aten::logical_not
+        lambda: OpenClipTextual("nllb-clip-large-siglip", (1, 77), pretrained="v1"), # ERROR (tinynn.converter.base) Unsupported ops: aten::logical_not
+        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="laion2b_e16"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="laion2b_e16"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="laion2b-s34b-b79k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="laion2b-s34b-b79k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-B-16", (1, 3, 224, 224), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-B-16", (1, 77), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-B-16", (1, 3, 224, 224), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-B-16", (1, 77), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-B-16-plus-240", (1, 3, 224, 224), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-B-16-plus-240", (1, 77), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-L-14", (1, 3, 224, 224), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-L-14", (1, 77), pretrained="laion400m_e31"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-L-14", (1, 3, 224, 224), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-L-14", (1, 77), pretrained="laion400m_e32"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-L-14", (1, 3, 224, 224), pretrained="laion2b-s32b-b82k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-L-14", (1, 77), pretrained="laion2b-s32b-b82k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-H-14", (1, 3, 224, 224), pretrained="laion2b-s32b-b79k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-H-14", (1, 77), pretrained="laion2b-s32b-b79k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("ViT-g-14", (1, 3, 224, 224), pretrained="laion2b-s12b-b42k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipTextual("ViT-g-14", (1, 77), pretrained="laion2b-s12b-b42k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("XLM-Roberta-Large-Vit-B-16Plus", (1, 3, 240, 240)), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("XLM-Roberta-Large-ViT-H-14", (1, 3, 224, 224), pretrained="frozen_laion5b_s13b_b90k"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("nllb-clip-base-siglip", (1, 3, 384, 384), pretrained="v1"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("nllb-clip-large-siglip", (1, 3, 384, 384), pretrained="v1"), # ERROR (tinynn.converter.base) Unsupported ops: aten::erf
+        lambda: OpenClipVisual("RN50", (1, 3, 224, 224), pretrained="yfcc15m"), # BatchNorm operation with mean/var output is not implemented
+        lambda: OpenClipTextual("RN50", (1, 77), pretrained="yfcc15m"), # BatchNorm operation with mean/var output is not implemented
+        lambda: OpenClipVisual("RN50", (1, 3, 224, 224), pretrained="cc12m"), # BatchNorm operation with mean/var output is not implemented
+        lambda: OpenClipTextual("RN50", (1, 77), pretrained="cc12m"), # BatchNorm operation with mean/var output is not implemented
+        lambda: MClipTextual("XLM-Roberta-Large-Vit-L-14", (1, 77)), # Expected normalized_shape to be at least 1-dimensional, i.e., containing at least one element, but got normalized_shape = []
+        lambda: MClipTextual("XLM-Roberta-Large-Vit-B-16Plus", (1, 77)), # Expected normalized_shape to be at least 1-dimensional, i.e., containing at least one element, but got normalized_shape = []
+        lambda: MClipTextual("LABSE-Vit-L-14", (1, 77)), # Expected normalized_shape to be at least 1-dimensional, i.e., containing at least one element, but got normalized_shape = []
+        lambda: OpenClipTextual("XLM-Roberta-Large-ViT-H-14", (1, 77), pretrained="frozen_laion5b_s13b_b90k"), # Expected normalized_shape to be at least 1-dimensional, i.e., containing at least one element, but got normalized_shape = []
+    ]
+    
+    succeeded: list[Callable[[], ExportBase]] = [
+        lambda: OpenClipVisual("ViT-B-32", (1, 3, 224, 224), pretrained="openai"),
+        lambda: OpenClipTextual("ViT-B-32", (1, 77), pretrained="openai"),
+        lambda: OpenClipVisual("ViT-B-16", (1, 3, 224, 224), pretrained="openai"),
+        lambda: OpenClipTextual("ViT-B-16", (1, 77), pretrained="openai"),
+        lambda: OpenClipVisual("ViT-L-14", (1, 3, 224, 224), pretrained="openai"),
+        lambda: OpenClipTextual("ViT-L-14", (1, 77), pretrained="openai"),
+        lambda: OpenClipVisual("ViT-L-14-336", (1, 3, 336, 336), pretrained="openai"),
+        lambda: OpenClipTextual("ViT-L-14-336", (1, 77), pretrained="openai"),
+        lambda: OpenClipVisual("RN50", (1, 3, 224, 224), pretrained="openai"),
+        lambda: OpenClipTextual("RN50", (1, 77), pretrained="openai"),
+        lambda: OpenClipTextual("ViT-H-14-quickgelu", (1, 77), pretrained="dfn5b"),
+        lambda: OpenClipTextual("ViT-H-14-378-quickgelu", (1, 77), pretrained="dfn5b"),
+        lambda: OpenClipVisual("XLM-Roberta-Large-Vit-L-14", (1, 3, 224, 224)),
+        lambda: OpenClipVisual("XLM-Roberta-Large-Vit-B-32", (1, 3, 224, 224)),
+        lambda: ArcFace("buffalo_s", (1, 3, 112, 112), optimization_level=3),
+        lambda: RetinaFace("buffalo_s", (1, 3, 640, 640), optimization_level=3),
+        lambda: ArcFace("buffalo_m", (1, 3, 112, 112), optimization_level=3),
+        lambda: RetinaFace("buffalo_m", (1, 3, 640, 640), optimization_level=3),
+        lambda: ArcFace("buffalo_l", (1, 3, 112, 112), optimization_level=3),
+        lambda: RetinaFace("buffalo_l", (1, 3, 640, 640), optimization_level=3),
+        lambda: ArcFace("antelopev2", (1, 3, 112, 112), optimization_level=3),
+        lambda: RetinaFace("antelopev2", (1, 3, 640, 640), optimization_level=3),
+    ]
+
+    models: list[Callable[[], ExportBase]] = [*failed, *succeeded]
+    for _model in succeeded:
+        model = _model().to(device)
+        try:
+            relative_path = os.path.join(model.task, "model.armnn")
+            output_path = os.path.join("output", model.model_name, relative_path)
+            model.to_armnn(output_path)
+            upload_file(path_or_fileobj=output_path, path_in_repo=relative_path, repo_id=model.repo_name)
+            if device == torch.device("cuda"):
+                model.half()
+                relative_path = os.path.join(model.task, "fp16", "model.armnn")
+                output_path = os.path.join("output", model.model_name, relative_path)
+                model.to_armnn(output_path)
+                upload_file(path_or_fileobj=output_path, path_in_repo=relative_path, repo_id=model.repo_name)
+            
+        except Exception as exc:
+            print(f"Failed to export {model.model_name} ({model.task}): {exc}")
+
+
+if __name__ == "__main__":
+    with torch.no_grad():
+        main()

machine-learning/export/ort/models/mclip.py

@@ -19,37 +19,44 @@ _MCLIP_TO_OPENCLIP = {
 }
 
 
+def forward(self: MultilingualCLIP, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        embs = self.transformer(input_ids, attention_mask)[0]
+        embs = (embs * attention_mask.unsqueeze(2)).sum(dim=1) / attention_mask.sum(dim=1)[:, None]
+        embs = self.LinearTransformation(embs)
+        return torch.nn.functional.normalize(embs, dim=-1)
+
+# unfortunately need to monkeypatch for tracing to work here
+# otherwise it hits the 2GiB protobuf serialization limit
+MultilingualCLIP.forward = forward
+
+
+def to_torchscript(model_name: str) -> torch.jit.ScriptModule:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        model = MultilingualCLIP.from_pretrained(model_name, cache_dir=tmpdir)
+
+        model.eval()
+        for param in model.parameters():
+            param.requires_grad_(False)
+        
+        return model
+
+
 def to_onnx(
     model_name: str,
     output_dir_visual: Path | str,
     output_dir_textual: Path | str,
 ) -> None:
     textual_path = get_model_path(output_dir_textual)
-    with tempfile.TemporaryDirectory() as tmpdir:
-        model = MultilingualCLIP.from_pretrained(model_name, cache_dir=tmpdir)
-        AutoTokenizer.from_pretrained(model_name).save_pretrained(output_dir_textual)
+    model = to_torchscript(model_name)
+    AutoTokenizer.from_pretrained(model_name).save_pretrained(output_dir_textual)
 
-        for param in model.parameters():
-            param.requires_grad_(False)
-
-        export_text_encoder(model, textual_path)
-        openclip_to_onnx(_MCLIP_TO_OPENCLIP[model_name], output_dir_visual)
-        optimize(textual_path)
+    _text_encoder_to_onnx(model, textual_path)
+    openclip_to_onnx(_MCLIP_TO_OPENCLIP[model_name], output_dir_visual)
+    optimize(textual_path)
 
 
-def export_text_encoder(model: MultilingualCLIP, output_path: Path | str) -> None:
+def _text_encoder_to_onnx(model: MultilingualCLIP, output_path: Path | str) -> None:
     output_path = Path(output_path)
-
-    def forward(self: MultilingualCLIP, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-        embs = self.transformer(input_ids, attention_mask)[0]
-        embs = (embs * attention_mask.unsqueeze(2)).sum(dim=1) / attention_mask.sum(dim=1)[:, None]
-        embs = self.LinearTransformation(embs)
-        return torch.nn.functional.normalize(embs, dim=-1)
-
-    # unfortunately need to monkeypatch for tracing to work here
-    # otherwise it hits the 2GiB protobuf serialization limit
-    MultilingualCLIP.forward = forward
-
     args = (torch.ones(1, 77, dtype=torch.int32), torch.ones(1, 77, dtype=torch.int32))
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", UserWarning)

machine-learning/export/ort/models/openclip.py

@@ -26,6 +26,17 @@ class OpenCLIPModelConfig:
         self.sequence_length = open_clip_cfg["text_cfg"]["context_length"]
 
 
+def to_torchscript(model_name: str) -> torch.jit.ScriptModule:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        model = MultilingualCLIP.from_pretrained(model_name, cache_dir=tmpdir)
+
+        model.eval()
+        for param in model.parameters():
+            param.requires_grad_(False)
+        
+        return model
+
+
 def to_onnx(
     model_cfg: OpenCLIPModelConfig,
     output_dir_visual: Path | str | None = None,
@@ -51,7 +62,7 @@ def to_onnx(
 
             save_config(open_clip.get_model_preprocess_cfg(model), output_dir_visual / "preprocess_cfg.json")
             save_config(text_vision_cfg, output_dir_visual.parent / "config.json")
-            export_image_encoder(model, model_cfg, visual_path)
+            _image_encoder_to_onnx(model, model_cfg, visual_path)
 
             optimize(visual_path)
 
@@ -61,11 +72,11 @@ def to_onnx(
 
             tokenizer_name = text_vision_cfg["text_cfg"].get("hf_tokenizer_name", "openai/clip-vit-base-patch32")
             AutoTokenizer.from_pretrained(tokenizer_name).save_pretrained(output_dir_textual)
-            export_text_encoder(model, model_cfg, textual_path)
+            _text_encoder_to_onnx(model, model_cfg, textual_path)
             optimize(textual_path)
 
 
-def export_image_encoder(model: open_clip.CLIP, model_cfg: OpenCLIPModelConfig, output_path: Path | str) -> None:
+def _image_encoder_to_onnx(model: open_clip.CLIP, model_cfg: OpenCLIPModelConfig, output_path: Path | str) -> None:
     output_path = Path(output_path)
 
     def encode_image(image: torch.Tensor) -> torch.Tensor:
@@ -89,7 +100,7 @@ def export_image_encoder(model: open_clip.CLIP, model_cfg: OpenCLIPModelConfig,
         )
 
 
-def export_text_encoder(model: open_clip.CLIP, model_cfg: OpenCLIPModelConfig, output_path: Path | str) -> None:
+def _text_encoder_to_onnx(model: open_clip.CLIP, model_cfg: OpenCLIPModelConfig, output_path: Path | str) -> None:
     output_path = Path(output_path)
 
     def encode_text(text: torch.Tensor) -> torch.Tensor: