DeepFillV2 (Inpaint) をONNX化して、量子化する

## 概要
過去にSI業務にてふと**消しゴムマジック**的なInpaintタスクに関する要望をいただき、調査した内容
結局、話は流れたが、その際、業務外で調べた、ONNX化や量子化周りの話が面白かったので、整理する
(消しゴムマジック的なアプリは個人的に作成して、今年のどこかで公開したい)

尚、以下記事と同一です

https://zenn.dev/t_shi/articles/ba5d63e4c96442

## 選定
- (元論文) [Free-Form Image Inpainting with Gated Convolution](https://arxiv.org/pdf/1806.03589)

- DeepFillv2
論文の著者のgithubリポジトリは https://github.com/JiahuiYu/generative_inpainting となる
但し、tensorflow実装となっており、contextual attentionに用いる、`tf.extract_image_patches`がONNX化に対応していなかった (pytorch版も同様)


苦戦した結果として、extract_image_patchesを愚直に実装している以下リポジトリに辿り着いた

https://github.com/csqiangwen/DeepFillv2_Pytorch

MITライセンスであったため、非商用目的として**fork**し、onnx化可能となるよう編集し、**今回公開した**

https://github.com/tominmin/DeepFillv2_Pytorch

尚、モデル学習は今回行わず、上記参考リポジトリ内URLにて共有される、checkpointから借用している
(今後、fine-tuning等、行っていきたい)

## 結論
onnx生成 (73MB) -> simplify (65MB) -> Uint8量子化 (17MB)まで成功し
最終モデルは、Inpaintとしての出力を得た
|元画像|mask|出力画像|
|--|--|--|
|![](https://storage.googleapis.com/zenn-user-upload/d40a596b178d-20240713.png)|![](https://storage.googleapis.com/zenn-user-upload/d9a2eb1e6a05-20240713.png)|![](https://storage.googleapis.com/zenn-user-upload/c780699521a0-20240713.png)|

[Netron](https://github.com/lutzroeder/netron)にて、モデル構成を出力した結果が以下で、軽量化の変遷が伺えると思う
|onnx|onnx (simplified) |onnx (uint8)|
|--|--|--|
|![](https://storage.googleapis.com/zenn-user-upload/6ec5e2ad425d-20240713.png)|![](https://storage.googleapis.com/zenn-user-upload/51e32f66eeb8-20240713.png)|![](https://storage.googleapis.com/zenn-user-upload/c4a45d45e327-20240713.png)



## 内容
タスクとしては、いわゆる**Inpaint**タスクとなる
> 「大きな欠陥領域(白抜き)を持つ画像を、もっともらしい合成画像で欠陥領域を埋めてあたかも欠陥がない画像にする」 [参照](https://qiita.com/ShozF/items/218a49841c47b2f360d7)

手順としては
1. pretrained_modelからmodel.loadし、torch.onnx.exportする
2. (1)のonnxファイルをonnxsimにて、simplifyする
3. (2)のonnxファイルをuint8量子化する
4. 動作確認

### 1. pretrained_modelからmodel.loadし、torch.onnx.exportする
少しだけ、OSS側を修正した
```python
# width = w.view(height, -1).data.shape[1]
width = w.view(height, int(w.numel() // height)).data.shape[1]
```
```python
# raw_w = raw_w.view(raw_int_bs[0], raw_int_bs[1], kernel, kernel, -1)
_ind = int(raw_w.numel() // (raw_int_bs[0] * raw_int_bs[1] * kernel * kernel))
raw_w = raw_w.view(raw_int_bs[0], raw_int_bs[1], kernel, kernel, _ind)
```
一部、`view`にて`reshape`処理を行っている箇所にて
[公式doc](https://pytorch.org/docs/0.3.0/tensors.html?highlight=view#torch.Tensor.view)にも、`the size -1 is inferred from other dimensions`と記載がある通り、よしなに推測して指定されるのかと思いきや、実行時エラーが発生した為
`numel` (= total number of elements)を利用し、計算して指定する必要があった

後は、OSS側の実装依存によりcpu演算が指定できていなかった為
今回の目的はonnx exportであることから、cpu演算の指定を直接行うよう軽い調整を行っている

他のDeepFillv2関連のgithub repositoryを調査・検証していた際に課題となっていた`extract_image_patches`に関しては、問題なくonnx exportされた
```python
generator = utils.create_generator(opt).eval()

# pretrained_modelは参考repository内URLからダウンロードする
model_name = "./(input path)/deepfillv2_WGAN_G_epoch40_batchsize4.pth"
...
pretrained_dict = torch.load(model_name, map_location=device)
generator.load_state_dict(pretrained_dict)

model = generator
torch.onnx.export(model, (dummy_image, dummy_mask), f="deepfillv2_mod.onnx")
```
上記手順にて、最初のonnxファイルが生成される

### 2. (1)のonnxファイルをonnxsimにて、simplifyする
onnx-simplifier (onnxsim)を用いて、モデル構造を機械的に最適化する

https://github.com/daquexian/onnx-simplifier

```python
from onnxsim import simplify

model = onnx.load(os.path.join(os.path.dirname(__file__), "deepfillv2_mod.onnx"))
model_simp, check = simplify(model)
...
onnx.save(model_simp, "deepfillv2_mod_simp.onnx")
```
上記手順にて、simplifyされたonnxファイルが生成される

### 3. (2)のonnxファイルをuint8量子化する
`onnxruntime.quantizationのquantize_dynamic`にて量子化する
`QInt8`でなく、`QUint8`を指定する必要がある点に注意する
```python
from onnxruntime.quantization import quantize_dynamic, QuantType
...
quantized_model = quantize_dynamic("deepfillv2_mod_simp.onnx", "deepfillv2_mod_simp_int8.onnx", weight_type=QuantType.QUInt8)
```
上記手順にて、Uint8量子化されたモデルが生成される

### 4. 動作確認
結果として、3段階の過程を得て、onnxファイルの容量は73MB => 65MB => 17MBまで軽量化された
17MB程度であれば、私の体感としては、m1 mac等であれば、8GBメモリ程度でも、wasm, SIMDを上手く利用すれば、なんとかweb上でも動作可能レベル。という認識である
![](https://storage.googleapis.com/zenn-user-upload/8cd976a5e27f-20240713.png)

実際に量子化済みのモデルにて、画像の欠陥補完ができているのか、確認した
|元画像|mask|出力画像|
|--|--|--|
|![](https://storage.googleapis.com/zenn-user-upload/d40a596b178d-20240713.png)|![](https://storage.googleapis.com/zenn-user-upload/d9a2eb1e6a05-20240713.png)|![](https://storage.googleapis.com/zenn-user-upload/c780699521a0-20240713.png)|

軽量化により、精度には懸念があるが・・補完としては機能し始めている
以降、調整していきたい

### 締め
LLM全盛の現在としては、隔世の感もあるが
軽量モデルの分野も、汎用性は無いが個別最適なソリューションとして地道に発展していくとは思う
分野としては、個人的に好きなので、ひっそりと趣味で調査していきたい

(以上)