SAM — the Segment Anything Model

SAM 01 - Overview & Architecture

Segment Anything Model (SAM) (Kirillov et al., Meta AI, 2023) is a promptable segmentation model. Given an image and a prompt (points, a box, a rough mask, or text via an add-on), it outputs one or more segmentation masks for the object(s) indicated by the prompt.

This series dissects SAM end to end using the Hugging Face transformers implementation (SamModel / SamProcessor) with real facebook/sam-vit-base weights when available, and a faithful synthetic fallback otherwise so every cell runs anywhere.

The three components

SAM factorizes segmentation into a heavy, prompt-independent image encoder run once per image, and a lightweight, prompt-dependent encoder+decoder that can be re-run many times per image (interactive use at ~50ms in a browser).

`` image (H,W,3) | +------------v-------------+ | Image Encoder | ViT-Det (MAE-pretrained) | 1024x1024x3 -> 64x64x256| run ONCE per image (heavy) +------------+-------------+ | image_embeddings (B,256,64,64) | prompts | +------------------+ points/box/mask -->+----------->| Prompt Encoder | (light) | +--------+---------+ | sparse (B,N,256) | dense (B,256,64,64) v v +--------------------------------------------+ | Mask Decoder | (light) | two-way transformer + upsampling + IoU | +--------------------+-----------------------+ masks (B,num_prompts,3,256,256) iou_scores (B,num_prompts,3) ``

The 3 in the mask output is SAM's multimask design: for a single ambiguous point it predicts 3 masks (e.g. part / subpart / whole) plus an IoU quality score for each, and you keep the best.

Load the real SAM config and print the architecture

When weights are present we read the actual config.json. Otherwise we fall back to the published sam-vit-base hyper-parameters so the numbers below are always correct for ViT-B.

=== Vision (image) encoder ===
  hidden_size      : 768
  num_hidden_layers: 12
  num_heads        : 12
  patch_size       : 16
  image_size       : 1024
  output_channels  : 256 (embedding dim C)
  global_attn_idx  : [2, 5, 8, 11]
=== Prompt encoder ===
  hidden_size          : 256
  image_embedding_size : 64 (64x64 grid)
=== Mask decoder ===
  hidden_size        : 256
  transformer layers : 2
  num_multimask_out  : 3

The key spatial numbers

1024 / 16 = 64, so the ViT produces a 64x64 grid of tokens; the neck projects the 768-dim tokens to a 256-dim embedding kept at 64x64. All downstream prompt/mask math happens in this 256-channel, 64x64 space.

Load the real model (optional, heavy)

This pulls the full SamModel into memory. It is used across later notebooks; here we just confirm it loads and report the parameter count per component.

Loaded SamModel on cuda
  vision_encoder params : 89,670,912
  prompt_encoder params : 6,220
  mask_decoder params   : 4,058,340
  TOTAL params          : 93,735,472

Why this factorization matters

• Amortized cost. The image encoder is ~93% of the parameters and all of the FLOPs; running it once lets the tiny decoder answer many prompts interactively.
• Promptable pre-training. Training SAM to resolve ambiguous prompts (via the 3-way multimask + min-over-masks loss) is what gives it zero-shot generality.
• Data engine. SAM was trained on SA-1B (11M images, 1.1B masks), bootstrapped with the model-in-the-loop (assisted -> semi-automatic -> fully automatic). NB05 covers this.

Next: NB02 opens up the image encoder tensor by tensor.

SAM 02 - Image Encoder (ViT-Det)

SAM's image encoder is a ViT pre-trained with MAE (masked auto-encoding) and adapted in the ViT-Det style for detection/segmentation: mostly windowed self-attention with a few global attention layers, followed by a small conv neck that produces the 256-channel image embedding.

Stage 1 - Patch embedding

A 16x16 strided conv turns the 1024x1024x3 image into a grid of 64x64 tokens of dimension C_vit (768 for ViT-B). A learned absolute position embedding of shape (1, 64, 64, C_vit) is added.

pixel_values : (1, 3, 1024, 1024)
after conv   : (1, 768, 64, 64) (B, C_VIT, G, G)
tokens (HWC) : (1, 64, 64, 768) (B, G, G, C_VIT)

Stage 2 - Transformer blocks: windowed vs global attention

Most blocks use windowed attention over 14x14 windows (cheap, local). A few blocks (indices [2,5,8,11] for ViT-B) use global attention across all 64*64 = 4096 tokens to mix information image-wide. This keeps the quadratic attention cost manageable while retaining global context.

Windowed attention on 14x14 windows costs O(n * w^4) vs global O(n^2) on n=4096 tokens - orders of magnitude cheaper.

global tokens        : 4096
windows (14x14)      : 25 of 196 tokens each
global attn FLOPs    : 2.577e+10
windowed attn FLOPs  : 1.475e+09
speedup (windowed)   : 17.5x

Stage 3 - The neck

After the last block, tokens (B, 64, 64, 768) are permuted to (B, 768, 64, 64) and passed through a small neck: 1x1 conv -> LayerNorm2d -> 3x3 conv -> LayerNorm2d, projecting 768 -> 256. The result is the image embedding (B, 256, 64, 64) consumed by the prompt encoder and mask decoder.

neck input       : (1, 768, 64, 64)
image_embeddings : (1, 256, 64, 64)

Run the REAL image encoder (if weights present)

We push a sample image through SamProcessor + sam_model.get_image_embeddings and confirm the (1,256,64,64) embedding, then visualize a few channels.

sample image: (480, 640, 3) from /mnt/mydata/huggingface/sample_cats.jpg

pixel_values     : (1, 3, 1024, 1024)

image_embeddings : (1, 256, 64, 64)

saved 02_image_embeddings.png; embedding stats: mean=0.012 std=0.155

Next: NB03 turns user prompts (points, boxes, masks) into the sparse and dense embeddings the decoder consumes.

SAM 03 - Prompt Encoder

The prompt encoder maps user prompts into the 256-dim embedding space:

• Points and boxes -> sparse embeddings (B, N, 256).
• A rough input mask -> dense embedding (B, 256, 64, 64), added directly to the image embedding.

Positions are encoded with a random Fourier / positional encoding so that geometry (where you clicked) is represented smoothly.

Points and boxes -> sparse embeddings

Each point gets a positional encoding of its (x,y) plus a learned embedding for its type: foreground point, background point, box top-left corner, box bottom-right corner. A box is thus encoded as two corner 'points'. SAM also appends learned output tokens (IoU token + 3 mask tokens) inside the decoder - covered in NB04.

point coords : (2, 2)
sparse embed : (2, 256) (N_points, C=256)

Dense (positional) grid

The prompt encoder also produces a dense positional encoding of the whole 64x64 grid, (256, 64, 64). When no mask prompt is supplied, a learned 'no-mask' embedding is broadcast to this shape instead; when a mask is supplied it is downsampled by convs to (256,64,64) and added to the image embedding.

dense positional encoding: (256, 64, 64) (C,64,64)
mask prompt in  : (1, 1, 256, 256)
dense mask embed: (1, 256, 64, 64) (C,64,64) added to image embedding

Run the REAL prompt encoder

We ask SamProcessor to prepare a point prompt, then call the model to get sparse & dense prompt embeddings, verifying shapes against the synthetic ones above.

real prompt encoder note: If points are provided, labels must also be provided.

Next: NB04 feeds image embedding + sparse/dense prompt embeddings into the two-way transformer mask decoder.

SAM 04 - Mask Decoder

The mask decoder is a lightweight two-way transformer that fuses the image embedding with the prompt tokens, then upsamples to masks and predicts an IoU quality score per mask.

Tokens going in

The decoder prepends learned output tokens to the sparse prompt tokens:

• 1 IoU token -> predicts mask quality scores.
• num_multimask_outputs + 1 = 4 mask tokens (3 for the multimask outputs + 1 used when a single mask is requested).

So the token sequence is [iou_token, mask_tokens(4), sparse_prompt(N)], of shape (B, 5+N, 256). The 'image side' is the flattened image embedding (B, 4096, 256) plus the dense positional encoding.

output+prompt tokens : 7 -> (B, 7 , 256)
image tokens         : 4096 -> (B, 4096, 256)
tokens       : (1, 7, 256)
image_tokens : (1, 4096, 256)

The two-way transformer

Each of the 2 layers does, in order:

1. Self-attention over the tokens.
2. Token -> image cross-attention (tokens read from the image).
3. MLP on tokens.
4. Image -> token cross-attention (image reads from tokens).

This bidirectional exchange lets prompts localize on the image and lets image features be modulated by the prompt. A final token->image attention updates the tokens once more.

token->image attn map: (1, 8, 7, 4096) (B,heads,n_tokens,4096)
updated tokens       : (1, 7, 256)
updated image tokens : (1, 4096, 256)

Upsampling to masks + IoU head

• The updated image tokens are reshaped to (B,256,64,64) and upscaled x4 by two transposed convs to (B,32,256,256).
• Each of the 3 mask tokens goes through a 3-layer MLP producing a 32-dim vector; a dot product with the upscaled features (per spatial location) yields the 256x256 mask logits.
• The IoU token goes through an MLP to predict the 3 quality scores.

Output: pred_masks (B,1,3,256,256), iou_scores (B,1,3).

upscaled features: (1, 32, 256, 256)
pred_masks       : (1, 3, 256, 256) (low-res 256x256)
iou_scores       : (1, 3)

Run the REAL decoder end-to-end

One forward pass of SamModel with a single point prompt returns 3 candidate masks + IoU scores. We visualize them ranked by predicted IoU.

pred_masks : (1, 1, 3, 256, 256)
iou_scores : (1, 1, 3)

saved 04_multimask.png

Next: NB05 covers how these outputs are trained - the losses and the SA-1B data engine.

SAM 05 - Training, Inputs/Outputs & Losses

SAM is trained for promptable segmentation: simulate an interactive prompting session on a ground-truth mask and learn to predict that mask. This notebook covers the exact inputs/outputs, the loss functions (focal + dice for masks, MSE for IoU), the multimask / ambiguity handling, and the SA-1B data engine.

Inputs and outputs during training

Inputs

• Image -> pixel_values (B,3,1024,1024).
• A simulated prompt for one target mask: a random point inside the GT mask, and/or the GT bounding box (optionally jittered), and/or a coarse mask from a previous iteration.

Outputs

• pred_masks (B,1,3,256,256) - 3 candidate masks.
• iou_scores (B,1,3) - predicted IoU for each candidate.

Targets

• The binary GT mask (downsampled to 256x256 for the loss).

Interactive simulation. Training runs up to ~11 iterations per mask: after each prediction, a new point is sampled from the error region (false-positive or false-negative) and fed back along with the previous logits as a mask prompt - mimicking a user refining a click.

The losses

The mask loss is a weighted sum of focal and dice loss:

$$\mathcal{L}_{mask} = 20\cdot\mathcal{L}_{focal} + 1\cdot\mathcal{L}_{dice}$$

(the paper uses a 20:1 focal:dice ratio). The IoU head is trained with MSE between the predicted IoU and the actual IoU of the predicted mask vs GT.

Ambiguity / multimask. Because a point can be ambiguous, SAM predicts 3 masks and backpropagates the mask loss only through the single lowest-loss mask (min-over-masks). The IoU head still learns to score all three.

mask 0: focal=0.003 dice=0.301 total=0.354 realIoU=1.000
mask 1: focal=0.064 dice=0.557 total=1.837 realIoU=0.478
mask 2: focal=0.466 dice=0.744 total=10.060 realIoU=0.147
min-over-masks backprops through mask 0
IoU-head MSE target for each mask = [1.0, 0.478, 0.147]

Visualize the loss landscape across the 3 masks

saved 05_losses.png

One synthetic training step (real model if present)

We run a single forward pass, compute the combined loss against a synthetic GT, and confirm gradients flow. This is illustrative - a full SA-1B run is out of scope, but the mechanics are identical.

combined loss: 3.914830207824707 | mask_decoder grad-norm sum: 1590.579

The SA-1B data engine

SAM's generality comes from SA-1B: 11M licensed images, 1.1B masks, built with the model in the loop in 3 stages:

1. Assisted-manual. Annotators click with SAM assisting; SAM is retrained on the growing data. (~4.3M masks)
2. Semi-automatic. SAM auto-fills confident masks; annotators label only what's missing, boosting diversity. (~10M masks)
3. Fully automatic. A 32x32 point grid prompts SAM per image; outputs are filtered by predicted IoU, stability, and NMS. This generated the vast majority (~1.1B masks).

The engine is a virtuous cycle: better model -> faster/higher-quality annotation -> more data -> better model.

saved 05_data_engine.png

Next: NB06 visualizes internal tensors - image embeddings, decoder attention, and mask logits.

SAM 06 - Layer Visualizations

We probe the model's internals on a real image: (1) image-embedding channel maps and their PCA, (2) the mask logits before thresholding, and (3) how the predicted mask sharpens across the sigmoid.

Image embedding: channels + PCA to RGB

The (256,64,64) embedding is hard to view directly. We show a few raw channels and a PCA projection of the 256-dim per-pixel features into 3 RGB components - a compact 'semantic colouring' of the scene.

embedding shape: (256, 64, 64)

saved 06_embedding_pca.png

Mask logits before/after sigmoid

For a point prompt we visualize the raw decoder logits and the probability map. The IoU-selected best mask usually has crisp, high-contrast logits.

saved 06_mask_logits.png; iou_scores = [0.844 0.813 0.824]

Next: NB07 runs full inference in point, box, and grid ('everything') modes on example images.

SAM 07 - Full Inference on Example Datasets

End-to-end inference in the three canonical modes:

1. Point prompt - click one object.
2. Box prompt - segment inside a bounding box.
3. 'Everything' / automatic - a grid of points to segment all objects (SAM's SamAutomaticMaskGenerator idea, reproduced with a point grid).

Uses cached sample images in /mnt/mydata/huggingface when present, else a synthetic scene.

image: (480, 640, 3) source: /mnt/mydata/huggingface/sample_cats.jpg

SAM ready on cuda

Mode 1 - Point prompt

saved 07_point_mode.png

Mode 2 - Box prompt

saved 07_box_mode.png

Mode 3 - 'Everything' via a point grid

We prompt SAM with an n x n grid of points (one forward pass per point batch), keep the best mask per point, filter by predicted IoU, and overlay the survivors with random colours - the essence of automatic mask generation.

saved 07_everything_mode.png; kept 35 masks

Summary

• One heavy image-encoder pass amortizes across many cheap prompts.
• Point/box/mask prompts all funnel through the same 256-dim, 64x64 embedding space and the same tiny two-way-transformer decoder.
• Multimask + IoU scoring resolves prompt ambiguity and enables the automatic 'everything' pipeline.

This completes the SAM deep-dive: architecture (NB01-04), training & losses (NB05), internals (NB06), and full inference (NB07).

SAM 08 - Fine-Tuning SAM (ViT-B) on Aerial Imagery

We fine-tune the smallest SAM checkpoint - facebook/sam-vit-base (the one downloaded in NB01) - on the harshinde/Aerial-Imagery dataset, using the Hugging Face transformers SamModel.

Strategy (the standard SAM recipe):

1. Freeze the ViT image encoder and the prompt encoder (they are already general); train only the mask decoder (~4M params, tiny).
2. For each GT instance, use its bounding box as the prompt and its binary mask as the target.
3. Loss = focal + dice on the decoder's predicted mask (upsampled to the label resolution), the same objective SAM was trained with.

Because the aerial masks are semantic (foreground vs background), each connected component becomes one 'instance' with a tight box prompt.

Step 1 - Get the Aerial-Imagery data

Reuse the copy the fastSAM notebook already downloaded if present, else snapshot it from the HF Hub, else synthesize a tiny stand-in so the notebook always runs.

AERIAL_DIR = /mnt/mydata/huggingface/aerial_imagery
image/mask pairs: 72

Step 2 - Build (image, box-prompt, GT-mask) training examples

For each image we take the largest few connected components of the foreground mask; each gives a box prompt (its bounding box, lightly jittered) and a target mask (that component). The SamProcessor handles resizing/normalization and produces input_boxes in SAM's 1024-space.

training instances: 277
train/val: 222 55

Step 3 - Load SAM ViT-B, freeze encoders, train only the mask decoder

We confirm the parameter split, then set requires_grad=False on the vision + prompt encoders. Only mask_decoder params get gradients.

model: facebook/sam-vit-base (smallest SAM)
total params    : 93,735,472
trainable params: 4,058,340 (mask decoder only, 4.3%)

Step 4 - Losses (focal + dice) and a helper to run one example

Step 5 - Baseline (pre-fine-tune) mask IoU on val

baseline val mask IoU (box-prompted, pre-FT): 0.4404

Step 6 - Fine-tune the mask decoder

A short training loop (few epochs, small batch) so the notebook completes quickly; raise SAM_FT_EPOCHS for a real run. Only the mask decoder updates.

epoch 1/4  mean loss=1.0530

epoch 2/4  mean loss=0.6862

epoch 3/4  mean loss=0.6128

epoch 4/4  mean loss=0.5803
fine-tuning done; steps: 888
saved 08_finetune_loss.png

Step 7 - Post-fine-tune IoU + before/after visualization

baseline val IoU : 0.4404
fine-tuned val IoU: 0.5058
delta IoU: 0.0654

saved 08_finetune_prediction.png

Notes

• We fine-tuned the smallest SAM (facebook/sam-vit-base) by training only the mask decoder - the ViT image encoder and prompt encoder stayed frozen. This is the memory/compute-efficient recipe: the decoder is a tiny fraction of the parameters, so it adapts quickly to the aerial domain.
• The prompt is each instance's bounding box (jittered during training for robustness) and the target is the GT mask; the objective is SAM's own focal + dice loss.
• For a stronger model: increase SAM_FT_EPOCHS, add point prompts / the interactive-refinement loop, and optionally unfreeze the prompt encoder. To save: model.save_pretrained(...).