ADD SAM2 (tflite)

[kyakuno]

edge-ai-torchで変換を検討。
https://github.com/facebookresearch/segment-anything-2
https://medium.com/axinc/ai-edge-torch%E3%81%A7pytorch%E3%81%8B%E3%82%89tflite%E3%81%AB%E5%A4%89%E6%8F%9B%E3%81%99%E3%82%8B-376be7dc5619
難しそうであれば、下記と同様に、Pytorch -> Kerasを検討。
https://github.com/tirthasheshpatel/segment_anything_keras

[kyakuno]

image_encoderをonnxには変換できるが、edge-ai-torchとCUDAでtfliteに変換しようとすると下記のエラーになる。

ValueError: Cannot view a tensor with shape torch.Size([1024, 4, 4, 288]) and strides (4608, 4, 1, 16) as a tensor with shape (16384, 288)!

While executing %view_38 : [num_users=1] = call_function[target=torch.ops.aten.view.default](args = (%view_37, [16384, 288]), kwargs = {})
Original traceback:
  File "/mnt/c/Users/kyakuno/Desktop/segment-anything-2/sam2/modeling/sam2_base.py", line 196, in forward
    backbone_out = self.forward_image(input_image)
  File "/mnt/c/Users/kyakuno/Desktop/segment-anything-2/sam2/modeling/sam2_base.py", line 485, in forward_image
    backbone_out = self.image_encoder(img_batch)
  File "/home/kyakuno/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/c/Users/kyakuno/Desktop/segment-anything-2/sam2/modeling/backbones/image_encoder.py", line 31, in forward
    features, pos = self.neck(self.trunk(sample))
  File "/home/kyakuno/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/c/Users/kyakuno/Desktop/segment-anything-2/sam2/modeling/backbones/hieradet.py", line 284, in forward
    x = blk(x)
  File "/home/kyakuno/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/c/Users/kyakuno/Desktop/segment-anything-2/sam2/modeling/backbones/hieradet.py", line 147, in forward
    x = self.attn(x)
  File "/home/kyakuno/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/mnt/c/Users/kyakuno/Desktop/segment-anything-2/sam2/modeling/backbones/hieradet.py", line 77, in forward
    x = self.proj(x)

[kyakuno]

segment_anything_kerasにSAM2対応のIssueはあるが未対応。
tirthasheshpatel/segment_anything_keras#4

[kyakuno]

edge-ai-torchをCPUモードで動かすとエラーが変わる。

<unknown>:0: error: failed while converting: 'main':
Some ops are not supported by the native TFLite runtime, you can enable TF kernels fallback using TF Select. See instructions: https://www.tensorflow.org/lite/guide/ops_select
TF Select ops: Relu6
Details:
        tf.Relu6(tensor<256x1xi64>) -> (tensor<256x1xi64>)
        tf.Relu6(tensor<256xi64>) -> (tensor<256xi64>)

[kyakuno]

tfliteのconverterには_ai_edge_converter_flagsでフラグを与えらえる。
https://github.com/google-ai-edge/ai-edge-torch/blob/main/docs/pytorch_converter/README.md

[kyakuno]

下記でFlexを有効にするとImageEncoderのエクスポート自体はできた。

            import ai_edge_torch
            import tensorflow as tf
            sample_inputs = (input_image,)
            tfl_converter_flags = {'target_spec': {'supported_ops': [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]}}
            edge_model = ai_edge_torch.convert(self.model, sample_inputs, _ai_edge_converter_flags=tfl_converter_flags)
            edge_model.export("image_encoder.tflite")

[kyakuno]

Relu6はFlexRelu6になる。

[kyakuno]

SegmentAnthing Int8の論文。ピークが2つある分布になるため、対処が必要と記載がある。
https://openaccess.thecvf.com/content/CVPR2024/papers/Lv_PTQ4SAM_Post-Training_Quantization_for_Segment_Anything_CVPR_2024_paper.pdf

[kyakuno]

edge-ai-torchでflexを有効にして量子化すると、MixedPrecisionのグラフになる。
Convはint8で、それ以外のオペレータはFloatで動く。

[kyakuno]

full int8 quantはまだサポートされていない気配がある。
google-ai-edge/ai-edge-torch#150

[kyakuno]

量子化対象のオペレータのリストが下記にある。
https://github.com/google-ai-edge/ai-edge-torch/blob/main/ai_edge_torch/quantize/pt2e_quantizer.py

[kyakuno]

PromptEncoderは下記のエラーになる。

RuntimeError: This model contains ops not capturable by Pytorch/XLA: aten::nonzero

エクスポートできない原因は、prompt_encoder.pyの_embed_pointsの下記のロジック。
ONNXだとWhereになる部分。

        point_embedding[labels == -1] = self.not_a_point_embed.weight
        point_embedding[labels == 0] += self.point_embeddings[0].weight
        point_embedding[labels == 1] += self.point_embeddings[1].weight
        point_embedding[labels == 2] += self.point_embeddings[2].weight
        point_embedding[labels == 3] += self.point_embeddings[3].weight

下記のようにするとエクスポートできるようになる。

        labels = labels.int()
        table = torch.zeros((5, self.point_embeddings[0].weight.shape[1]))
        table[0] = self.not_a_point_embed.weight
        table[1] = self.point_embeddings[0].weight
        table[2] = self.point_embeddings[1].weight
        table[3] = self.point_embeddings[2].weight
        table[4] = self.point_embeddings[3].weight
        for i in range(labels.shape[0]):
            point_embedding[i] = point_embedding[i] + table[labels[i] + 1]

[kyakuno]

ImageEncoderをtfliteで推論してみる。ImageEncoderは正常にexportできている。

[kyakuno]

MaskDecoderは下記のエラーになる。

  File "/home/kyakuno/.local/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 221, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: call_function ConstantVariable(int: 512) [ConstantVariable()] {}

from user code:
   File "/mnt/c/Users/kyakuno/Desktop/segment-anything-2/sam2/modeling/sam/mask_decoder.py", line 137, in forward
    masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks(
  File "/mnt/c/Users/kyakuno/Desktop/segment-anything-2/sam2/modeling/sam/mask_decoder.py", line 198, in predict_masks
    sparse_prompt_embeddings.size(0), -1, -1

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

I0000 00:00:1724378905.868815    9151 cpu_client.cc:470] TfrtCpuClient destroyed.

[kyakuno]

size -> shapeに置き換えるとここはpassする。

[kyakuno]

Floatモデルはすべてエクスポートできた
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/mask_decoder_hiera_l.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/prompt_encoder_sparse_hiera_l.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/image_encoder_hiera_l.tflite

[kyakuno]

量子化でキャリブレーションしようとすると下記のエラーになる。

torch.histogram: input tensor and hist tensor should have the same dtype, but got input long int and hist float

pytorch/pytorch#74420

[kyakuno]

量子化のフロー
https://pytorch.org/tutorials/prototype/quantization_in_pytorch_2_0_export_tutorial.html

[kyakuno]

グラフの入力が問題ではなく、グラフの途中でint64のテンソルが出てきて対応できなくなっている。

[kyakuno]

torch/ao/quantization/observer.pyのreset_histogramで下記のコードを追加すると通る。

    def reset_histogram(self, x: torch.Tensor, min_val: torch.Tensor, max_val: torch.Tensor) -> None:
        self.min_val.resize_(min_val.shape)
        self.min_val.copy_(min_val)
        self.max_val.resize_(max_val.shape)
        self.max_val.copy_(max_val)
        assert (
            min_val.numel() == 1 and max_val.numel() == 1
        ), "histogram min/max values must be scalar."
        if x.dtype != torch.float32: # 追加
            x = x.float() # 追加
        torch.histc(
            x, self.bins, min=min_val, max=max_val, out=self.histogram  # type: ignore[arg-type]
        )

[kyakuno]

int8版のImageEncoderの出力。魂はあっていそう。

[kyakuno]

使用したバージョン
torch 2.4.0
ai-edge-torch 0.2.0

[kyakuno]

PromptEncoderはlabelがint64で量子化できない。

[kyakuno]

PromptEncoderは演算量が少ないのでfloatで動かしてもいい気はする。

[kyakuno]

MaskDecoderの量子化結果。ImageEncoderはfloat。

[kyakuno]

量子化モデル。ただし、キャリブレーションはtruckでしか行っていない。

https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/mask_decoder_hiera_l_int8.tflite
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/image_encoder_hiera_l_int8.tflite

[kyakuno]

ai-edge-torchがLLMを対象にしているためか、テンソルは結構、floatになっている。
量子化ツールに手を入れないといけない感じはする。

[kyakuno]

ImageEncoderのAttentionのところは、weightはint8になっていて、floatにして行列積を行っている。

[kyakuno]

ImageEncoderの出力にposを含めると、量子化でエラーが発生する。
is_dynamic=Trueにして、DynamicQuantizationにすると通るが、演算は全てFloatになる。

[kyakuno]

ImageEncoderのDynamicQuantizationの出力。出力は綺麗。

[kyakuno]

is_dynamic = False（後処理もモデルに含むことでposを出力に含めない）
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/image_encoder_hiera_l_int8_is_dynamc_false.tflite

is_dynamic = True（posを出力に含める）
https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/image_encoder_hiera_l_int8_is_dynamic_true.tflite

[kyakuno]

ImageEncoder、PromptEncoder、MaskDecoderは正常にtflite (float)に変換できた。

MemoryAttentionの変換は下記で行う。
axinc-ai/ailia-models#1514

[kyakuno]

MaskDecoderをdynamic shapeにしようとすると、下記のエラーになる。

 File "/home/kyakuno/.local/lib/python3.10/site-packages/torch_xla/experimental/unbounded_dynamism_export.py", line 115, in decompose_dynamic_shape_select
    assert symbolic_dims[
AssertionError: Selected dim cannot be symbolic

tfliteはStatic Shapeのみ対応にする。