🦊

Tensorflowjsのモデルをnpmパッケージにバンドルする方法

2021/10/05に公開

JavaScript

TypeScript

tech

はじめに

これまで様々な機械学習モデルをWebWorkderで動かすnpmパッケージを作成してきたのですが、npm install後にtensorflowjsのモデルをコピーする必要があって、インストールが少し面倒でした。

どうにかwebpackでバンドルしてリリースできないかなぁと調査したら、できるようになったので、ここでやり方を紹介したいと思います。

なお、他者が作成したモデルをバンドルする場合はライセンスには十分気をつけるようにしてください。また、モデルをバンドルしてしまうと、後から新しいモデルに入れ替えるというのが難しくなりますので、その点もご了承の上行うようにしてください。

前提

今回はwebpack5系を前提とします。4系では試してないです。

モデルはBiseNetv2-Tensorflowをtensorflowjs化したものを使います。なお、このモデルをtensorflowjs化する方法は、こちらの記事をご参照ください。

npmパッケージの作成

まずは、準備としてtensorflowjsのモデルをnpmパッケージのルートに配置してください。
ここでは、resourcesというフォルダの下に置くことにします。

$ ls resources/bisenetv2-celebamask/
group1-shard1of3.bin  group1-shard2of3.bin  group1-shard3of3.bin  model.json

次にこれらのモデルファイルをバンドルするためのルールをwebpack.config.jsに設定します。
webpack5からはurl-loaderやfile-loaderが不要になり、設定が簡単になりました。
JSONはasset/source, バイナリはasset/inlineとしておくのが良いでしょう。

        rules: [
            { test: /\.ts$/, loader: 'ts-loader' },
            { test:/resources\/.*\.bin/, type:"asset/inline"},
            { test:/resources\/.*\.json/, type:"asset/source"}            
        ],

npmパッケージのコードは次のようになります。

import * as tf from '@tensorflow/tfjs';

import modelJson from "../resources/bisenetv2-celebamask/model.json" // <------ (1)
import modelWeight1 from "../resources/bisenetv2-celebamask/group1-shard1of3.bin" // <------ (2)
import modelWeight2 from "../resources/bisenetv2-celebamask/group1-shard2of3.bin" // <------ (2)
import modelWeight3 from "../resources/bisenetv2-celebamask/group1-shard3of3.bin" // <------ (2)

export class BiseNetV2{
    model: tf.GraphModel | null = null
    canvas = document.createElement("canvas")

    init = async () => {
        const modelJson2 = new File([modelJson], "model.json", {type: "application/json"}) // <------ (3)
        const b1 = Buffer.from(modelWeight1.split(',')[1], 'base64')  // <------ (4)
        const modelWeights1 = new File([b1], "group1-shard1of3.bin")  // <------ (4)
        const b2 = Buffer.from(modelWeight2.split(',')[1], 'base64')  // <------ (4)
        const modelWeights2 = new File([b2], "group1-shard2of3.bin")  // <------ (4)
        const b3 = Buffer.from(modelWeight3.split(',')[1], 'base64')  // <------ (4)
        const modelWeights3 = new File([b3], "group1-shard3of3.bin")  // <------ (4)
        this.model = await tf.loadGraphModel(tf.io.browserFiles([modelJson2, modelWeights1, modelWeights2, modelWeights3]))   // <------ (5)
    }
    predict = async (targetCanvas: HTMLCanvasElement, processWidth: number, processHeight: number): Promise<number[][]> => {
        <snip...>
    }

}

(1)でJSONファイルをインポートしています。
(2)でバイナリファイルをインポートしています。
(3)でJSONファイルのデータをファイルとして読み込みます。
(4)でバイナリファイルのデータをファイルとして読み込みます。
(5)でファイルを指定してtensorflowjsのモデルをロードしています。
なお、今回はGraphModelをロードしていますが、LayersModelも同じ方法でロードできます。

以上です。やり方がわかればとても簡単ですね。

npmパッケージを使用するアプリケーションの作成

普通に使うだけなら、パッケージをインポートして、newするだけで使用できます。
しかし、通常tensorflowjsのモデルをバンドルしたパッケージは大きなサイズになりがちです。
このため、使用するときにだけloadするのが良いでしょう。

もしcreate-react-appを用いてアプリケーションを作成しているのであれば、
コード分割をしておくと、使用するときだけロードすることができます。
次のようにimport関数をするようにしておくと、create-react-appはそのモジュールを異なるchunkに分割してくれます。

    const mod = await import('bisenetv2-js')
    const lib = new mod.BiseNetV2()

リポジトリ

ここに記載した内容は下記のリポジトリに格納されています。

Readmeに従ってコマンドを実行してみてください。
コード量も少ないので、動かしながらで１０分位で理解できると思います。

さいごに

ここまで、tensorflowjsのモデルのバンドルの仕方を紹介してきました。
しかし、最初に述べたとおり、tensorflowjsのモデルのバンドルは諸刃の剣です。
パッケージのインストールは簡単になるかもしれませんが、次のようなデメリットもあります。

モデルの入れ替えが難しくなる
パッケージのサイズが大きくなる
改変不可のモデルはバンドルできない（たぶん）

バンドルしたほうが良いかどうかは、皆様がしっかりと判断していただくようお願いします。