| """Export SAM 3's vision graph once for both prompts: the Perception Encoder read of a photo feeds |
| the detector's neck (words) and the tracker's neck (clicks and boxes). |
| |
| Outputs: features_288/144/72 (as vision.onnx, for the detector) and image_embeddings.0/1/2 (as |
| onnx-community's sam3-tracker encoder writes them, for its prompt decoder). |
| python export_shared.py <weights dir> <out dir> |
| """ |
| import sys, time |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
| from transformers import Sam3VideoModel |
|
|
| from sam3_graphs import IMAGE_SIZE |
|
|
|
|
| class Sam3SharedVision(nn.Module): |
| def __init__(self, video): |
| super().__init__() |
| self.vision_encoder = video.detector_model.vision_encoder |
| self.tracker_neck = video.tracker_neck |
| self.conv_s0 = video.tracker_model.mask_decoder.conv_s0 |
| self.conv_s1 = video.tracker_model.mask_decoder.conv_s1 |
| self.no_memory_embedding = video.tracker_model.no_memory_embedding |
|
|
| def forward(self, pixel_values: torch.Tensor): |
| out = self.vision_encoder(pixel_values) |
| levels = out.fpn_hidden_states |
| hidden = out.last_hidden_state |
| b, side = hidden.shape[0], IMAGE_SIZE // 14 |
| spatial = hidden.view(b, side, side, -1).permute(0, 3, 1, 2) |
| tracked, _ = self.tracker_neck(spatial) |
| e0 = self.conv_s0(tracked[0]) |
| e1 = self.conv_s1(tracked[1]) |
| e2 = tracked[2] + self.no_memory_embedding.reshape(1, -1, 1, 1) |
| return levels[0], levels[1], levels[2], e0, e1, e2 |
|
|
|
|
| if __name__ == "__main__": |
| weights, out = Path(sys.argv[1]), Path(sys.argv[2]) |
| out.mkdir(parents=True, exist_ok=True) |
| t = time.time() |
| video = Sam3VideoModel.from_pretrained(weights, attn_implementation="eager", torch_dtype=torch.float32).eval() |
| print(f"loaded in {time.time() - t:.0f} s", flush=True) |
| module = Sam3SharedVision(video).eval() |
| pixels = torch.zeros(1, 3, IMAGE_SIZE, IMAGE_SIZE) |
| with torch.no_grad(): |
| shapes = [tuple(x.shape) for x in module(pixels)] |
| print("outputs", shapes, flush=True) |
| t = time.time() |
| with torch.no_grad(): |
| torch.onnx.export(module, (pixels,), str(out / "vision.onnx"), input_names=["pixel_values"], |
| output_names=["features_288", "features_144", "features_72", |
| "image_embeddings.0", "image_embeddings.1", "image_embeddings.2"], |
| opset_version=18, dynamo=True, external_data=True, optimize=True) |
| print(f"exported in {time.time() - t:.0f} s", flush=True) |
|
|
