File size: 5,596 Bytes
179df43 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | # SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
"""
PyTorch reference pipeline for SAM3 inference.
Used to generate reference outputs for comparison with ttnn implementation.
"""
import torch
import torch.nn.functional as F
from typing import Dict, List, Optional, Tuple
def preprocess_image(image: torch.Tensor, target_size: int = 1008) -> torch.Tensor:
"""
Preprocess image for SAM3 inference.
Args:
image: (B, C, H, W) or (C, H, W) tensor in [0, 255] or [0, 1] range
target_size: target image size (SAM3 uses 1008)
Returns:
Preprocessed image tensor (B, 3, target_size, target_size) normalized to [-1, 1]
"""
if image.ndim == 3:
image = image.unsqueeze(0)
# Normalize to [0, 1] if needed
if image.max() > 1.0:
image = image.float() / 255.0
# Resize
if image.shape[-2:] != (target_size, target_size):
image = F.interpolate(
image, size=(target_size, target_size), mode="bilinear", align_corners=False
)
# Normalize to [-1, 1] (SAM3 uses mean=0.5, std=0.5)
image = (image - 0.5) / 0.5
return image
def run_vit_backbone(model, pixel_values: torch.Tensor) -> torch.Tensor:
"""
Run just the ViT backbone portion of SAM3.
Args:
model: SAM3 ViT backbone (sam3_model.backbone.visual.trunk)
pixel_values: (B, 3, 1008, 1008) preprocessed image
Returns:
ViT features (B, 72, 72, 1024)
"""
with torch.no_grad():
features = model(pixel_values)
return features
def run_vit_single_block(block, hidden_states: torch.Tensor) -> torch.Tensor:
"""
Run a single ViT block.
Args:
block: ViT Block module
hidden_states: (B, H, W, C) tensor
Returns:
(B, H, W, C) tensor
"""
with torch.no_grad():
output = block(hidden_states)
return output
def run_neck(neck_model, pixel_values: torch.Tensor) -> Dict:
"""
Run the full visual backbone (ViT + FPN neck).
Args:
neck_model: SAM3 visual backbone (sam3_model.backbone.visual)
pixel_values: (B, 3, 1008, 1008) preprocessed image
Returns:
Dict with 'backbone_fpn' (list of feature maps) and 'vision_pos_enc' (position encodings)
"""
with torch.no_grad():
output = neck_model(pixel_values)
return output
def run_text_encoder(text_encoder, text_prompts: List[str]) -> torch.Tensor:
"""
Run the text encoder on text prompts.
Args:
text_encoder: SAM3 text encoder
text_prompts: list of text strings
Returns:
Text embeddings tensor
"""
with torch.no_grad():
text_features = text_encoder(text_prompts)
return text_features
def extract_vit_block_params(block) -> Dict[str, torch.Tensor]:
"""
Extract all parameters from a ViT Block for ttnn preprocessing.
Args:
block: ViT Block module (sam3.model.vitdet.Block)
Returns:
Dict mapping parameter names to tensors
"""
params = {}
# LayerNorm 1
params["norm1_weight"] = block.norm1.weight.data.clone()
params["norm1_bias"] = block.norm1.bias.data.clone()
# Attention QKV (fused)
params["attn_qkv_weight"] = block.attn.qkv.weight.data.clone()
params["attn_qkv_bias"] = block.attn.qkv.bias.data.clone()
# Attention output projection
params["attn_proj_weight"] = block.attn.proj.weight.data.clone()
params["attn_proj_bias"] = block.attn.proj.bias.data.clone()
# RoPE frequencies (if present)
if hasattr(block.attn, "freqs_cis") and block.attn.freqs_cis is not None:
params["attn_freqs_cis"] = block.attn.freqs_cis.clone()
# Relative position biases (if present)
if hasattr(block.attn, "rel_pos_h") and block.attn.rel_pos_h is not None:
params["attn_rel_pos_h"] = block.attn.rel_pos_h.data.clone()
params["attn_rel_pos_w"] = block.attn.rel_pos_w.data.clone()
# LayerNorm 2
params["norm2_weight"] = block.norm2.weight.data.clone()
params["norm2_bias"] = block.norm2.bias.data.clone()
# MLP
params["mlp_fc1_weight"] = block.mlp.fc1.weight.data.clone()
params["mlp_fc1_bias"] = block.mlp.fc1.bias.data.clone()
params["mlp_fc2_weight"] = block.mlp.fc2.weight.data.clone()
params["mlp_fc2_bias"] = block.mlp.fc2.bias.data.clone()
# Window size
params["window_size"] = block.window_size
return params
def extract_vit_backbone_params(vit_model) -> Dict:
"""
Extract all parameters from the ViT backbone.
Args:
vit_model: SAM3 ViT model (sam3.model.vitdet.ViT)
Returns:
Dict with all model parameters organized by component
"""
params = {}
# Patch embedding
params["patch_embed"] = {
"weight": vit_model.patch_embed.proj.weight.data.clone(),
"bias": (
vit_model.patch_embed.proj.bias.data.clone()
if vit_model.patch_embed.proj.bias is not None
else None
),
}
# Position embeddings
if hasattr(vit_model, "pos_embed") and vit_model.pos_embed is not None:
params["pos_embed"] = vit_model.pos_embed.data.clone()
# Pre-LayerNorm (if present)
if hasattr(vit_model, "ln_pre") and vit_model.ln_pre is not None:
params["ln_pre_weight"] = vit_model.ln_pre.weight.data.clone()
params["ln_pre_bias"] = vit_model.ln_pre.bias.data.clone()
# Blocks
params["blocks"] = []
for block in vit_model.blocks:
params["blocks"].append(extract_vit_block_params(block))
return params
|