[WIP] Export to tflite int8

[kyakuno]

高速化のため、FULLY_INTEGER_QUANTIZATIONのモデルの構築の検討。

Windowsだとedge-ai-torchの依存ライブラリが入らないので、WSLが必要。

 python3 export_image_predictor.py --image_size 512 --framework tflite --accuracy int8

現段階では、下記の変換エラーが発生する。

   output = self._dispatch_impl(func, types, args, kwargs)
  File "/home/kyakuno/.local/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1757, in _dispatch_impl
    r = func(*args, **kwargs)
  File "/home/kyakuno/.local/lib/python3.10/site-packages/torch/_ops.py", line 667, in __call__
    return self_._op(*args, **kwargs)
  File "/home/kyakuno/.local/lib/python3.10/site-packages/torch/ao/quantization/fx/_decomposed.py", line 81, in quantize_per_tensor_meta
    assert input.dtype == torch.float32, f"Expecting input to have dtype torch.float32, but got dtype: {input.dtype}"
torch._dynamo.exc.TorchRuntimeError: Failed running call_function quantized_decomposed.quantize_per_tensor.default(*(FakeTensor(..., size=(), dtype=torch.int64), 0.007843137718737125, -128, -128, 127, torch.int8), **{}):
Expecting input to have dtype torch.float32, but got dtype: torch.int64

from user code:
   File "<eval_with_key>.12", line 1083, in forward
    quantize_per_tensor_default_261 = torch.ops.quantized_decomposed.quantize_per_tensor.default(_tensor_constant_9, 0.007843137718737125, -128, -128, 127, torch.int8);  _tensor_constant_9 = None

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

[kyakuno]

エラーの_tensor_constant_9はMul -> div -> Powに繋がっている。
position_encoding.pyでdim_tを計算しているところで、torch.arrangeが定数としてグラフの入力になってしまい、量子化に失敗している。
dim_tを事前計算にすると、変換が行えるようになる。

[kyakuno]

キャリブレーション画像をtruck.jpgとした場合の出力。グラフの精度は足りそう。

int8

float

[kyakuno]

TensorWise Quantization版
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/image_encoder_hiera_t_2.1_512.torch_tw_int8.tflite

[kyakuno]

ChannelWise Quantization版
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/image_encoder_hiera_t_2.1_512.torch_cw_int8.tflite

[kyakuno]

Image Encoderはレンジのばらつきはあまりないようで、1枚の画像でキャリブレーションしても正常な絵が出力される。

[kyakuno]

推論サンプル
axinc-ai/ailia-models-tflite#95

[kyakuno]

mask decoderは下記のエラーになる。

ValueError: num_scales must be 1 for per-layer quantization, or 256 for per-axis quantization, but got 128.Tensor 201 has invalid quantization parameters.num_scales must be 1 for per-layer quantization, or 256 for per-axis quantization, but got 128.Tensor 217 has invalid quantization parameters.num_scales must be 1 for per-layer quantization, or 256 for per-axis quantization, but got 128.Tensor 270 has invalid quantization parameters.

[kyakuno]

is_per_channel = Falseにした場合も、Convはper channelで量子化される。
FCで新しい規格のis_per_channelを使うかどうかの判定みたい。

[kyakuno]

Memory Attention Per Channel
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/memory_attention_hiera_t_2.1_512.torch_cw_int8.tflite
Mask Decoder Per Tensor
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/mask_decoder_hiera_t_2.1_512.torch_tw_int8.tflite

[kyakuno]

キャリブレーション用のデータセットを作成する機能を足さないといけない。

[kyakuno]

下記の例では、torchを使用せずに、tensorflowの方で量子化している。
google-ai-edge/ai-edge-torch#345

[kyakuno]

tensorflowで量子化すると、layernormもint8になる。
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/image_encoder_hiera_t_2.1_512.tf_int8.tflite

torch (layernormがfloatになる)

tensorflow (layernormがint8になる)

[kyakuno]

全てのモデルにキャリブレーション機能を追加。とりあえず、COCOデータセットの100枚でキャリブレーションをかける。

[kyakuno]

キャリブレーションを実施。image_encoder、mask_decoder、prompt_encoderはCOCOデータセットの100枚、それ以外はbedroomの200枚を使用。
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/image_encoder_hiera_t_2.1_512.int8.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/mask_decoder_hiera_t_2.1_512.int8.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/memory_attention_hiera_t_2.1_512.int8.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/obj_ptr_tpos_proj_hiera_t_2.1_512.int8.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/memory_encoder_hiera_t_2.1_512.int8.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/prompt_encoder_hiera_t_2.1_512.int8.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/mlp_hiera_t_2.1_512.int8.tflite

[kyakuno]

テスト方法

./download_tflite_models.sh
python3 export_image_predictor.py --framework tflite --accuracy int8 --mode import --image_size 512
python3 export_video_predictor.py --framework tflite --accuracy int8 --mode import --image_size 512

出力

output/*.png

[kyakuno]

tensorflowではなくtorchで量子化し、layernormをFloatで実行したモデル。キャリブレーション条件は同一。
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/layernorm_float/image_encoder_hiera_t_2.1_512.int8.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2.1/layernorm_float/mask_decoder_hiera_t_2.1_512.int8.tflite