cd checkpoints && \
./download_ckpts.sh && \
cd ..onnx
torch 2.2.1
onnx 1.16.2
tflite
torch 2.4.0
ai-edge-torch 0.2.0
tf-nightly 2.18.0.dev20240905
onnx
python3 export_image_predictor.py --framework onnx
python3 export_video_predictor.py --framework onnx
tflite
export PJRT_DEVICE=CPU
python3 export_image_predictor.py --framework tflite
python3 export_video_predictor.py --framework tflite
onnx
python3 export_image_predictor.py --framework onnx --mode import
python3 export_video_predictor.py --framework onnx --mode import
tflite not supported inference only yet.
Replacing the complex tensor of RotaryEnc with matmul. To test this behavior, you can also run it with torch.
python3 export_video_predictor.py --framework torch
The deliverables will be stored below.
output/*
model/*
You can also download it from the following.
- https://storage.googleapis.com/ailia-models/segment-anything-2/image_encoder_hiera_t.onnx
- https://storage.googleapis.com/ailia-models/segment-anything-2/prompt_encoder_hiera_t.onnx
- https://storage.googleapis.com/ailia-models/segment-anything-2/mask_decoder_hiera_t.onnx
- https://storage.googleapis.com/ailia-models/segment-anything-2/memory_encoder_hiera_t.onnx
- https://storage.googleapis.com/ailia-models/segment-anything-2/mlp_hiera_t.onnx
- https://storage.googleapis.com/ailia-models/segment-anything-2/memory_attention_hiera_t.onnx (6dim matmul, batch = N)
- https://storage.googleapis.com/ailia-models/segment-anything-2/memory_attention_hiera_t.opt.onnx (4dim matmul, batch = 1)
- https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/image_encoder_hiera_t.tflite
- https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/prompt_encoder_hiera_t.tflite
- https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/mask_decoder_hiera_t.tflite
- https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/mlp_hiera_t.tflite
- https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/memory_encoder_hiera_t.tflite
- https://storage.googleapis.com/ailia-models-tflite/segment-anything-2/memory_attention_hiera_t.tflite (4dim matmul, batch = 1, num_maskmem = 1)
The memory attention in tflite does not support dynamic shapes, so num_maskmem and max_obj_ptrs_in_encoder need to be fixed to 1.