From 0829b88776951e2ba1c5bc41b0fd488b4b92c5ef Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Mon, 24 Nov 2025 20:59:52 -0800 Subject: [PATCH] Add files via upload --- sam3/perflib/__init__.py | 8 + sam3/perflib/associate_det_trk.py | 137 ++++++ sam3/perflib/compile.py | 99 +++++ sam3/perflib/connected_components.py | 84 ++++ sam3/perflib/fa3.py | 27 ++ sam3/perflib/masks_ops.py | 69 +++ sam3/perflib/nms.py | 91 ++++ sam3/perflib/tests/assets/masks.tiff | Bin 0 -> 352484 bytes sam3/perflib/tests/tests.py | 59 +++ sam3/perflib/triton/connected_components.py | 468 ++++++++++++++++++++ sam3/perflib/triton/nms.py | 124 ++++++ sam3/sam/__init__.py | 4 + sam3/sam/common.py | 39 ++ sam3/sam/mask_decoder.py | 319 +++++++++++++ sam3/sam/prompt_encoder.py | 243 ++++++++++ sam3/sam/rope.py | 161 +++++++ sam3/sam/transformer.py | 358 +++++++++++++++ 17 files changed, 2290 insertions(+) create mode 100644 sam3/perflib/__init__.py create mode 100644 sam3/perflib/associate_det_trk.py create mode 100644 sam3/perflib/compile.py create mode 100644 sam3/perflib/connected_components.py create mode 100644 sam3/perflib/fa3.py create mode 100644 sam3/perflib/masks_ops.py create mode 100644 sam3/perflib/nms.py create mode 100644 sam3/perflib/tests/assets/masks.tiff create mode 100644 sam3/perflib/tests/tests.py create mode 100644 sam3/perflib/triton/connected_components.py create mode 100644 sam3/perflib/triton/nms.py create mode 100644 sam3/sam/__init__.py create mode 100644 sam3/sam/common.py create mode 100644 sam3/sam/mask_decoder.py create mode 100644 sam3/sam/prompt_encoder.py create mode 100644 sam3/sam/rope.py create mode 100644 sam3/sam/transformer.py diff --git a/sam3/perflib/__init__.py b/sam3/perflib/__init__.py new file mode 100644 index 0000000..5c3823b --- /dev/null +++ b/sam3/perflib/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import os + +is_enabled = False +if os.getenv("USE_PERFLIB", "1") == "1": + # print("Enabled the use of perflib.\n", end="") + is_enabled = True diff --git a/sam3/perflib/associate_det_trk.py b/sam3/perflib/associate_det_trk.py new file mode 100644 index 0000000..508ae81 --- /dev/null +++ b/sam3/perflib/associate_det_trk.py @@ -0,0 +1,137 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from collections import defaultdict + +import torch +import torch.nn.functional as F +from sam3.perflib.masks_ops import mask_iou +from scipy.optimize import linear_sum_assignment + + +def associate_det_trk( + det_masks, + track_masks, + iou_threshold=0.5, + iou_threshold_trk=0.5, + det_scores=None, + new_det_thresh=0.0, +): + """ + Optimized implementation of detection <-> track association that minimizes DtoH syncs. + + Args: + det_masks: (N, H, W) tensor of predicted masks + track_masks: (M, H, W) tensor of track masks + + Returns: + new_det_indices: list of indices in det_masks considered 'new' + unmatched_trk_indices: list of indices in track_masks considered 'unmatched' + """ + with torch.autograd.profiler.record_function("perflib: associate_det_trk"): + assert isinstance(det_masks, torch.Tensor), "det_masks should be a tensor" + assert isinstance(track_masks, torch.Tensor), "track_masks should be a tensor" + if det_masks.size(0) == 0 or track_masks.size(0) == 0: + return list(range(det_masks.size(0))), [], {}, {} # all detections are new + + if list(det_masks.shape[-2:]) != list(track_masks.shape[-2:]): + # resize to the smaller size to save GPU memory + if torch.numel(det_masks[-2:]) < torch.numel(track_masks[-2:]): + track_masks = ( + F.interpolate( + track_masks.unsqueeze(1).float(), + size=det_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ).squeeze(1) + > 0 + ) + else: + # resize detections to track size + det_masks = ( + F.interpolate( + det_masks.unsqueeze(1).float(), + size=track_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ).squeeze(1) + > 0 + ) + + det_masks = det_masks > 0 + track_masks = track_masks > 0 + + iou = mask_iou(det_masks, track_masks) # (N, M) + igeit = iou >= iou_threshold + igeit_any_dim_1 = igeit.any(dim=1) + igeit_trk = iou >= iou_threshold_trk + + iou_list = iou.cpu().numpy().tolist() + igeit_list = igeit.cpu().numpy().tolist() + igeit_any_dim_1_list = igeit_any_dim_1.cpu().numpy().tolist() + igeit_trk_list = igeit_trk.cpu().numpy().tolist() + + det_scores_list = ( + det_scores + if det_scores is None + else det_scores.cpu().float().numpy().tolist() + ) + + # Hungarian matching for tracks (one-to-one: each track matches at most one detection) + # For detections: allow many tracks to match to the same detection (many-to-one) + + # If either is empty, return all detections as new + if det_masks.size(0) == 0 or track_masks.size(0) == 0: + return list(range(det_masks.size(0))), [], {} + + # Hungarian matching: maximize IoU for tracks + cost_matrix = 1 - iou.cpu().numpy() # Hungarian solves for minimum cost + row_ind, col_ind = linear_sum_assignment(cost_matrix) + + def branchy_hungarian_better_uses_the_cpu( + cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks + ): + matched_trk = set() + matched_det = set() + matched_det_scores = {} # track index -> [det_score, det_score * iou] det score of matched detection mask + for d, t in zip(row_ind, col_ind): + matched_det_scores[t] = [ + det_scores_list[d], + det_scores_list[d] * iou_list[d][t], + ] + if igeit_trk_list[d][t]: + matched_trk.add(t) + matched_det.add(d) + + # Tracks not matched by Hungarian assignment above threshold are unmatched + unmatched_trk_indices = [ + t for t in range(track_masks.size(0)) if t not in matched_trk + ] + + # For detections: allow many tracks to match to the same detection (many-to-one) + # So, a detection is 'new' if it does not match any track above threshold + assert track_masks.size(0) == igeit.size( + 1 + ) # Needed for loop optimizaiton below + new_det_indices = [] + for d in range(det_masks.size(0)): + if not igeit_any_dim_1_list[d]: + if det_scores is not None and det_scores[d] >= new_det_thresh: + new_det_indices.append(d) + + # for each detection, which tracks it matched to (above threshold) + det_to_matched_trk = defaultdict(list) + for d in range(det_masks.size(0)): + for t in range(track_masks.size(0)): + if igeit_list[d][t]: + det_to_matched_trk[d].append(t) + + return ( + new_det_indices, + unmatched_trk_indices, + det_to_matched_trk, + matched_det_scores, + ) + + return (branchy_hungarian_better_uses_the_cpu)( + cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks + ) diff --git a/sam3/perflib/compile.py b/sam3/perflib/compile.py new file mode 100644 index 0000000..f427aa7 --- /dev/null +++ b/sam3/perflib/compile.py @@ -0,0 +1,99 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +def recursive_fn_factory(fn): + def recursive_fn(b): + if isinstance(b, dict): + return {k: recursive_fn(b[k]) for k in b} + if isinstance(b, list): + return [recursive_fn(t) for t in b] + if isinstance(b, tuple): + return tuple(recursive_fn(t) for t in b) + if isinstance(b, torch.Tensor): + return fn(b) + # Yes, writing out an explicit white list of + # trivial types is tedious, but so are bugs that + # come from not applying fn, when expected to have + # applied it. + if b is None: + return b + trivial_types = [bool, int] + for t in trivial_types: + if isinstance(b, t): + return b + raise TypeError(f"Unexpected type {type(b)}") + + return recursive_fn + + +recursive_contiguous = recursive_fn_factory(lambda x: x.contiguous()) +recursive_clone = recursive_fn_factory(torch.clone) + + +def compile_wrapper( + fn, *, mode="max-autotune", fullgraph=True, dynamic=False, name=None +): + compiled_fn = torch.compile(fn, mode=mode, fullgraph=fullgraph, dynamic=dynamic) + + def compiled_fn_wrapper(*args, **kwargs): + with torch.autograd.profiler.record_function( + f"compiled {fn}" if name is None else name + ): + cont_args = recursive_contiguous(args) + cont_kwargs = recursive_contiguous(kwargs) + result = compiled_fn(*cont_args, **cont_kwargs) + cloned_result = recursive_clone(result) + return cloned_result + + return compiled_fn_wrapper + + +def shape_logging_wrapper(fn, keep_kwargs, enable_logging=False): + """ + Wraps a function and prints the shapes of all tensor inputs. + Only prints when a new combination of shapes is seen. + Thread-safe. + + Args: + fn: Function to wrap + enable_logging: Boolean flag to enable/disable logging + """ + seen_shapes = set() + + def get_shape(obj): + if isinstance(obj, torch.Tensor): + return obj.shape + elif isinstance(obj, (list, tuple)): + if len(obj) > 1: + return tuple(get_shape(x) for x in obj) + return get_shape(obj[0]) + elif isinstance(obj, dict): + return tuple(sorted((k, get_shape(v)) for k, v in obj.items())) + else: + return type(obj).__name__ + + def wrapper(*args, **kwargs): + shapes = tuple(get_shape(arg) for arg in args) + tuple( + (k, get_shape(v)) + for k, v in kwargs.items() + if isinstance(v, (torch.Tensor, list)) + and (len(keep_kwargs) > 0 and k in keep_kwargs) + ) + if shapes not in seen_shapes: + seen_shapes.add(shapes) + if enable_logging: + print(f"[ShapeLogger] New input shapes for {fn.__qualname__}: {shapes}") + return fn(*args, **kwargs) + + # Allow toggling the flag at runtime + wrapper.enable_logging = enable_logging + + def set_logging(enabled=False): + nonlocal enable_logging + enable_logging = enabled + wrapper.enable_logging = enable_logging + + wrapper.set_logging = set_logging + return wrapper diff --git a/sam3/perflib/connected_components.py b/sam3/perflib/connected_components.py new file mode 100644 index 0000000..c96932a --- /dev/null +++ b/sam3/perflib/connected_components.py @@ -0,0 +1,84 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import logging + +import torch + +try: + from cc_torch import get_connected_components + + HAS_CC_TORCH = True +except ImportError: + logging.debug( + "cc_torch not found. Consider installing for better performance. Command line:" + " pip install git+https://github.com/ronghanghu/cc_torch.git" + ) + HAS_CC_TORCH = False + + +def connected_components_cpu_single(values: torch.Tensor): + assert values.dim() == 2 + from skimage.measure import label + + labels, num = label(values.cpu().numpy(), return_num=True) + labels = torch.from_numpy(labels) + counts = torch.zeros_like(labels) + for i in range(1, num + 1): + cur_mask = labels == i + cur_count = cur_mask.sum() + counts[cur_mask] = cur_count + return labels, counts + + +def connected_components_cpu(input_tensor: torch.Tensor): + out_shape = input_tensor.shape + if input_tensor.dim() == 4 and input_tensor.shape[1] == 1: + input_tensor = input_tensor.squeeze(1) + else: + assert ( + input_tensor.dim() == 3 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + batch_size = input_tensor.shape[0] + labels_list = [] + counts_list = [] + for b in range(batch_size): + labels, counts = connected_components_cpu_single(input_tensor[b]) + labels_list.append(labels) + counts_list.append(counts) + labels_tensor = torch.stack(labels_list, dim=0).to(input_tensor.device) + counts_tensor = torch.stack(counts_list, dim=0).to(input_tensor.device) + return labels_tensor.view(out_shape), counts_tensor.view(out_shape) + + +def connected_components(input_tensor: torch.Tensor): + """ + Computes connected components labeling on a batch of 2D tensors, using the best available backend. + + Args: + input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Both tensors have the same shape as input_tensor. + - A tensor with dense labels. Background is 0. + - A tensor with the size of the connected component for each pixel. + """ + if input_tensor.dim() == 3: + input_tensor = input_tensor.unsqueeze(1) + + assert ( + input_tensor.dim() == 4 and input_tensor.shape[1] == 1 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + if input_tensor.is_cuda: + if HAS_CC_TORCH: + return get_connected_components(input_tensor.to(torch.uint8)) + else: + # triton fallback + from sam3.perflib.triton.connected_components import ( + connected_components_triton, + ) + + return connected_components_triton(input_tensor) + + # CPU fallback + return connected_components_cpu(input_tensor) diff --git a/sam3/perflib/fa3.py b/sam3/perflib/fa3.py new file mode 100644 index 0000000..8f8c9bd --- /dev/null +++ b/sam3/perflib/fa3.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +@torch.library.custom_op("flash::flash_attn_func", mutates_args=()) +def flash_attn_func_op( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor +) -> torch.Tensor: + from flash_attn_interface import flash_attn_func as fa3 + + return fa3(q, k, v) + + +def flash_attn_func(q, k, v): + dtype = torch.float8_e4m3fn + return flash_attn_func_op(q.to(dtype), k.to(dtype), v.to(dtype)).to(q.dtype) + + +@flash_attn_func_op.register_fake +def _(q, k, v, **kwargs): + # two outputs: + # 1. output: (batch, seq_len, num_heads, head_dim) + # 2. softmax_lse: (batch, num_heads, seq_len) with dtype=torch.float32 + # output needs to be bfloat16, not float8! + meta_q = torch.empty_like(q, dtype=torch.bfloat16).contiguous() + return meta_q diff --git a/sam3/perflib/masks_ops.py b/sam3/perflib/masks_ops.py new file mode 100644 index 0000000..48299d5 --- /dev/null +++ b/sam3/perflib/masks_ops.py @@ -0,0 +1,69 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +def masks_to_boxes(masks: torch.Tensor, obj_ids: list[int]): + with torch.autograd.profiler.record_function("perflib: masks_to_boxes"): + # Sanity check based on callsite for replacement + assert masks.shape[0] == len(obj_ids) + assert masks.dim() == 3 + + # Based on torchvision masks_to_boxes + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device, dtype=torch.float) + + N, H, W = masks.shape + device = masks.device + y = torch.arange(H, device=device).view(1, H) + x = torch.arange(W, device=device).view(1, W) + + masks_with_obj = masks != 0 # N, H, W + masks_with_obj_x = masks_with_obj.amax( + dim=1 + ) # N, H (which columns have objects) + masks_with_obj_y = masks_with_obj.amax(dim=2) # N, W (which rows have objects) + masks_without_obj_x = ~masks_with_obj_x + masks_without_obj_y = ~masks_with_obj_y + + bounding_boxes_0 = torch.amin( + (masks_without_obj_x * W) + (masks_with_obj_x * x), dim=1 + ) + bounding_boxes_1 = torch.amin( + (masks_without_obj_y * H) + (masks_with_obj_y * y), dim=1 + ) + bounding_boxes_2 = torch.amax(masks_with_obj_x * x, dim=1) + bounding_boxes_3 = torch.amax(masks_with_obj_y * y, dim=1) + + bounding_boxes = torch.stack( + [bounding_boxes_0, bounding_boxes_1, bounding_boxes_2, bounding_boxes_3], + dim=1, + ).to(dtype=torch.float) + assert bounding_boxes.shape == (N, 4) + assert bounding_boxes.device == masks.device + assert bounding_boxes.dtype == torch.float + return bounding_boxes + + +def mask_iou(pred_masks: torch.Tensor, gt_masks: torch.Tensor) -> torch.Tensor: + """ + Compute the IoU (Intersection over Union) between predicted masks and ground truth masks. + Args: + - pred_masks: (N, H, W) bool Tensor, containing binary predicted segmentation masks + - gt_masks: (M, H, W) bool Tensor, containing binary ground truth segmentation masks + Returns: + - ious: (N, M) float Tensor, containing IoUs for each pair of predicted and ground truth masks + """ + assert pred_masks.dtype == gt_masks.dtype == torch.bool + N, H, W = pred_masks.shape + M, _, _ = gt_masks.shape + + # Flatten masks: (N, 1, H*W) and (1, M, H*W) + pred_flat = pred_masks.view(N, 1, H * W) + gt_flat = gt_masks.view(1, M, H * W) + + # Compute intersection and union: (N, M) + intersection = (pred_flat & gt_flat).sum(dim=2).float() + union = (pred_flat | gt_flat).sum(dim=2).float() + ious = intersection / union.clamp(min=1) + return ious # shape: (N, M) diff --git a/sam3/perflib/nms.py b/sam3/perflib/nms.py new file mode 100644 index 0000000..b3efc59 --- /dev/null +++ b/sam3/perflib/nms.py @@ -0,0 +1,91 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging + +import numpy as np +import torch + +from sam3.perflib.masks_ops import mask_iou + + +try: + from torch_generic_nms import generic_nms as generic_nms_cuda + + GENERIC_NMS_AVAILABLE = True +except ImportError: + logging.debug( + "Falling back to triton or CPU mask NMS implementation -- please install `torch_generic_nms` via\n\t" + 'pip uninstall -y torch_generic_nms; TORCH_CUDA_ARCH_LIST="8.0 9.0" pip install git+https://github.com/ronghanghu/torch_generic_nms' + ) + GENERIC_NMS_AVAILABLE = False + + +def nms_masks( + pred_probs: torch.Tensor, + pred_masks: torch.Tensor, + prob_threshold: float, + iou_threshold: float, +) -> torch.Tensor: + """ + Args: + - pred_probs: (num_det,) float Tensor, containing the score (probability) of each detection + - pred_masks: (num_det, H_mask, W_mask) float Tensor, containing the binary segmentation mask of each detection + - prob_threshold: float, score threshold to prefilter detections (NMS is performed on detections above threshold) + - iou_threshold: float, mask IoU threshold for NMS + + Returns: + - keep: (num_det,) bool Tensor, indicating whether each detection is kept after score thresholding + NMS + """ + # prefilter the detections with prob_threshold ("valid" are those above prob_threshold) + is_valid = pred_probs > prob_threshold # (num_det,) + probs = pred_probs[is_valid] # (num_valid,) + masks_binary = pred_masks[is_valid] > 0 # (num_valid, H_mask, W_mask) + if probs.numel() == 0: + return is_valid # no valid detection, return empty keep mask + + ious = mask_iou(masks_binary, masks_binary) # (num_valid, num_valid) + kept_inds = generic_nms(ious, probs, iou_threshold) + + # valid_inds are the indices among `probs` of valid detections before NMS (or -1 for invalid) + valid_inds = torch.where(is_valid, is_valid.cumsum(dim=0) - 1, -1) # (num_det,) + keep = torch.isin(valid_inds, kept_inds) # (num_det,) + return keep + + +def generic_nms( + ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5 +) -> torch.Tensor: + """A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix.""" + + assert ious.dim() == 2 and ious.size(0) == ious.size(1) + assert scores.dim() == 1 and scores.size(0) == ious.size(0) + + if ious.is_cuda: + if GENERIC_NMS_AVAILABLE: + return generic_nms_cuda(ious, scores, iou_threshold, use_iou_matrix=True) + else: + from sam3.perflib.triton.nms import nms_triton + + return nms_triton(ious, scores, iou_threshold) + + return generic_nms_cpu(ious, scores, iou_threshold) + + +def generic_nms_cpu( + ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5 +) -> torch.Tensor: + """ + A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix. (CPU implementation + based on https://github.com/jwyang/faster-rcnn.pytorch/blob/master/lib/model/nms/nms_cpu.py) + """ + ious_np = ious.float().detach().cpu().numpy() + scores_np = scores.float().detach().cpu().numpy() + order = scores_np.argsort()[::-1] + kept_inds = [] + while order.size > 0: + i = order.item(0) + kept_inds.append(i) + inds = np.where(ious_np[i, order[1:]] <= iou_threshold)[0] + order = order[inds + 1] + + return torch.tensor(kept_inds, dtype=torch.int64, device=scores.device) diff --git a/sam3/perflib/tests/assets/masks.tiff b/sam3/perflib/tests/assets/masks.tiff new file mode 100644 index 0000000000000000000000000000000000000000..7a8efc6dd0ee958b3ebf2733a0bfb32d39449478 GIT binary patch literal 352484 zcmeI&&#Gis6$ju`_x8QXAOw1l1d~a~D2NCKL2%$apidy}pFtV}4gxw6@d11R(TPJJ z!6%51;5yZd~*>UKZguDfr>?{02}?~I@Qb9_JF@7->CuUYecbLI!H)_gq- zA8zgr_kOe=|7Ses(`N4v#`E79Kl|_a9$wrHfBfWkzxnl7Uwrc8kAME7?|uB)XFuHi z^5Re1FZZu)rS1FoU%q_#^2^_T@%^v<&@VJ2FjHXDzGj|vecRh`yYcmvIHRHm>9!xt zIPc1~9*tKf;E8lqT6iX2g%+NTS0G^Vd_lbI4GiNY4_G|kgm|eNSUlgPc!?V*-%jAl zez!cm9R1V3=}*hkuk?I+KA--nn|`C`)AKp~X?gmU)t}GtZ+a%Yp5A@28z-bo{J#$) z(}aMFVXkpd^?Y#y*HB;b&Y*)!UbV_u)$^TyLRTSQ_|h{ST==H7jysWXt#X#X_h`W7 zuUYZL2MJfK<#Dd=5;WXFX;m+-$7uN=WFhzhw9_O+^bgEAICbZSf#$XmmaE*2Xn7l zWq%y&uws?^=3aWJJ|4`yYL)$Qtiy^`>YID%q560*_o`L)$FUA8R;h39rHAU{!Q87> z*&oL`tXQSK*_R&5j|Q`ETIYZ$*5T{xlsEg4tt^ewm|{Zdz-Y2i{=tb5MMfa{jE-0dd?cx<~VDWtWcvTyCh*vaV@qA6Zk_{}LuZvf(fyMK^7Oz?ZAH*vb zuz0>B@k%wYc)p|Y3N^5Jz7z4vG_ZKSlksXa@JzfC0gLB58?QhEi{}gCrEg&Id||xY z4V(}!b-?2JCdEtKz~cEP#!K74;`t`WOWVL~|N23_1PBlyK!5-N0t5&UAV7cs0RjXF z5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk z1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs z0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZ zfB*pk1PBlya0LR-+Se7-vup`%xZQZ!<5xcILAvb+E7!E(^+)5?33wu2m4GMX)d+Yd zUWI^XqHz8i?2Ff=IEchG2r#?#m^zW&kxl8Wze0n}Fd0w7+ zt>@G8`P9$cCHGZ-K8FbbFY&s%wbk>*4ZOORTvkxnz{`%izSTj%>nr%uip~bS^ssAO ze2_K>@~+5wB_8{?I3VDWsrcy$}Nk5~35%&UAw10LcPeL{=p zYvPq`VDWriym}3MEnc}Fp|9sF7Vts5VozxCd`IGyYGCnvN8{CL;E8x;euSRNS0v!c zctxJj;`z?RE78E>`Oe16-@qVV`X8ZDzT^SJc*&p8;`t`TOWnZY`6k6n-N38;i&ljZ zAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C7 z2oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N z0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+ z009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBng zVuAZBu5E$R?(b^{ES_(SSGs}4^X=jlZea0z`*>9wc!*atVDWrSypjzpp0A5nuz|(% zy%w)p10TdI7O;4}Bk@W#uz0?s@d`Dtc)kBctVTkI~%V+ z1B>Si;-zn3@qA&t+zp%%FLl7;`6k6n+`!`bCdNzKz~cEP$IIEk3*w~=SUlfF@e($$ zc)kncrE6gEd>6+{*T4n-v6>?Y5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs z0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZ zfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&U zAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FpSCymiT| z-#jL459;Qgt7a50YR`gp_%$n-s zVTwe3k#VSIMZY!0tf@X8rbyHm8HZ|C^jlNRn(E_WibQ>paj0fRzct0IsXiX2NYob@ zhiX>zTT{%M>f>RGM17HQsAfgKHN~u{J|3n>)E60tYF6}HQ_Py`<6(+KeUWjfW<|d> z#jL459;Qgt7a50YR`gp_%$n-sVTwe3k#VSIMZY!0tf@X8rbyHm8HZ|C^jlNRn(E_W zibQ>paj0fRzct0IsXiX2NYob@hiX>zTT{%M>f>RGM17HQsAfgKHN~u{J|3n>)E60t zYF6}HQ_Py`<6(+KeUWjfW<|d>#jL459;Qgt7a50YR`gp_%$n-sVTwe3k#VSIMZY!0 ztf@X8rbyHm8HZ|C^jlNRn(E_WibQ>paj0fRzct0IsXiX2NYob@hiX>zTT{%M>f>RG zM17HQsAfgKHN~u{J|3n>)E60tYF6}HQ_Py`<6(+KeUWjfW<|d>#jL459;Qgt7a50Y zR`gp_%$n-sVTwe3k#VSIMZY!0tf@X8rbyHm8HZ|C^jlNRn(E_WibQ>paj0fRzct0I zsXiX2NYob@hiX>zTT{%M>f>RGM17HQsAfgK)rFOn z(~NI6x4X~Z9zVa`kGJdYoALAg`r%*0@Lu!z`^`BYyjt`1FnqYV+nxEB7k9%)`|;n$ z@u$t+e~k0(J?(p;D*wD$tjeGM;?b)7+09~Ae*1V;KEGY8%D?{d(JcSHT$Mlh)uUDU Rw`Yr4j^8-lch2^^=zrO2=;8nX literal 0 HcmV?d00001 diff --git a/sam3/perflib/tests/tests.py b/sam3/perflib/tests/tests.py new file mode 100644 index 0000000..0fb88ad --- /dev/null +++ b/sam3/perflib/tests/tests.py @@ -0,0 +1,59 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import os + +import numpy as np +import pytest +import torch +from PIL import Image +from sam3.perflib.masks_ops import masks_to_boxes + + +class TestMasksToBoxes: + def test_masks_box(self): + def masks_box_check(masks, expected, atol=1e-4): + out = masks_to_boxes(masks, [1 for _ in range(masks.shape[0])]) + assert out.dtype == torch.float + print("out: ", out) + print("expected: ", expected) + torch.testing.assert_close( + out, expected, rtol=0.0, check_dtype=True, atol=atol + ) + + # Check for int type boxes. + def _get_image(): + assets_directory = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "assets" + ) + mask_path = os.path.join(assets_directory, "masks.tiff") + image = Image.open(mask_path) + return image + + def _create_masks(image, masks): + for index in range(image.n_frames): + image.seek(index) + frame = np.array(image) + masks[index] = torch.tensor(frame) + + return masks + + expected = torch.tensor( + [ + [127, 2, 165, 40], + [2, 50, 44, 92], + [56, 63, 98, 100], + [139, 68, 175, 104], + [160, 112, 198, 145], + [49, 138, 99, 182], + [108, 148, 152, 213], + ], + dtype=torch.float, + ) + + image = _get_image() + for dtype in [torch.float16, torch.float32, torch.float64]: + masks = torch.zeros( + (image.n_frames, image.height, image.width), dtype=dtype + ) + masks = _create_masks(image, masks) + masks_box_check(masks, expected) diff --git a/sam3/perflib/triton/connected_components.py b/sam3/perflib/triton/connected_components.py new file mode 100644 index 0000000..253ca9d --- /dev/null +++ b/sam3/perflib/triton/connected_components.py @@ -0,0 +1,468 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import math + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _any_combine(a, b): + return a | b + + +@triton.jit +def tl_any(a, dim=0): + return tl.reduce(a, dim, _any_combine) + + +# ============================================================================== +# ## Phase 1: Initialization Kernel +# ============================================================================== +# Each foreground pixel (value > 0) gets a unique label equal to its +# linear index. Background pixels (value == 0) get a sentinel label of -1. +# Note that the indexing is done across batch boundaries for simplicity +# (i.e., the first pixel of image 1 gets label H*W, etc.) + + +@triton.jit +def _init_labels_kernel( + input_ptr, labels_ptr, numel: tl.constexpr, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + input_values = tl.load(input_ptr + offsets, mask=mask, other=0) + + indices = tl.where((input_values != 0), offsets, -1) + tl.store(labels_ptr + offsets, indices, mask=mask) + + +# ============================================================================== +# ## Phase 2: Local merging +# ============================================================================== +# Each pixel tries to merge with its 8-connected neighbors (up, down, left, right) +# if they have the same value. This is done using a disjoint-set union operation. + + +@triton.jit +def find(labels_ptr, indices, mask): + current_pids = indices + + # 'is_done' tracks lanes that have finished their work. + # A lane is initially "done" if it's not active (mask is False). + is_done = ~mask + + # Loop as long as there is at least one lane that is NOT done. + while tl_any(~is_done): + # The work_mask is for lanes that are still active and seeking their root. + work_mask = ~is_done + parents = tl.load(labels_ptr + current_pids, mask=work_mask, other=-1) + # A lane is now done if its parent is itself (it's a root) + # or if it hits a -1 sentinel (a safe exit condition). + is_root = parents == current_pids + is_sentinel = parents == -1 + is_done |= is_root | is_sentinel + + # For lanes that are not yet done, update their pid to their parent to continue traversal. + current_pids = tl.where(is_done, current_pids, parents) + # We could add the following line to do path compression, but experimentally it's slower + # tl.atomic_min(labels_ptr + indices, current_pids, mask=mask) + return current_pids + + +@triton.jit +def union(labels_ptr, a, b, process_mask): + # This function implements a disjoint-set union + # As an invariant, we use the fact that the roots have the lower id. That helps parallelization + # However, that is not sufficient by itself. Suppose two threads want to do union(0,2) and union(1,2) at the same time + # Then if we do a naive atomic_min, 0 and 1 will compete to be the new parent of 2 and min(0, 1) will win. + # However, 1 still needs to be merged with the new {0, 2} component. + # To ensure that merge is also done, we need to detect whether the merge was successful, and if not retry until it is + + current_a = a + current_b = b + + final_root = a + # A mask to track which lanes have successfully completed their union. + done_mask = ~process_mask # tl.zeros_like(a) == 1 # Init with all False + + while tl_any(~done_mask): + # Define the mask for lanes that still need work in this iteration + work_mask = process_mask & ~done_mask + + # Find the roots for the current a and b values in the active lanes + root_a = find(labels_ptr, current_a, work_mask) + tl.debug_barrier() + root_b = find(labels_ptr, current_b, work_mask) + + # 7. Merge logic + # If roots are already the same, the sets are already merged. Mark as done. + are_equal = root_a == root_b + final_root = tl.where(are_equal & work_mask & ~done_mask, root_a, final_root) + done_mask |= are_equal & work_mask + + # Define masks for the two merge cases (a < b or b < a) + a_is_smaller = root_a < root_b + + # Case 1: root_a < root_b. Attempt to set parent[root_b] = root_a + merge_mask_a_smaller = work_mask & a_is_smaller & ~are_equal + ptr_b = labels_ptr + root_b + old_val_b = tl.atomic_min(ptr_b, root_a, mask=merge_mask_a_smaller) + + # A lane is done if its atomic op was successful (old value was what we expected) + success_b = old_val_b == root_b + final_root = tl.where(success_b & work_mask & ~done_mask, root_a, final_root) + done_mask |= success_b & merge_mask_a_smaller + + # *** Crucial Retry Logic *** + # If the update failed (old_val_b != root_b), another thread interfered. + # We update `current_b` to this new root (`old_val_b`) and will retry in the next loop iteration. + current_b = tl.where(success_b | ~merge_mask_a_smaller, current_b, old_val_b) + + # Case 2: root_b < root_a. Attempt to set parent[root_a] = root_b + merge_mask_b_smaller = work_mask & ~a_is_smaller & ~are_equal + ptr_a = labels_ptr + root_a + old_val_a = tl.atomic_min(ptr_a, root_b, mask=merge_mask_b_smaller) + + success_a = old_val_a == root_a + final_root = tl.where(success_a & work_mask & ~done_mask, root_b, final_root) + done_mask |= success_a & merge_mask_b_smaller + + # *** Crucial Retry Logic *** + # Similarly, update `current_a` if the atomic operation failed. + current_a = tl.where(success_a | ~merge_mask_b_smaller, current_a, old_val_a) + + return final_root + + +@triton.jit +def _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + dx: tl.constexpr, + dy: tl.constexpr, +): + # Helper functions to compute merge with a specific neighbor offset (dx, dy) + + neighbor_h = offsets_h + dy + neighbor_w = offsets_w + dx + # Proper bounds checking: all four bounds must be satisfied + mask_n = ( + mask_2d + & (neighbor_h[:, None] >= 0) + & (neighbor_h[:, None] < H) + & (neighbor_w[None, :] >= 0) + & (neighbor_w[None, :] < W) + ) + + offsets_neighbor = neighbor_h[:, None] * W + neighbor_w[None, :] + neighbor_values = tl.load( + input_ptr + base_offset + offsets_neighbor, mask=mask_n, other=-1 + ) + + mask_n = tl.ravel(mask_n) + neighbor_labels = tl.load( + labels_ptr + tl.ravel(base_offset + offsets_neighbor), mask=mask_n, other=-1 + ) + + to_merge = ( + mask_n & (neighbor_labels != -1) & tl.ravel(current_values == neighbor_values) + ) + valid_write = valid_current & to_merge + + # returns new parents for the pixels that were merged (otherwise keeps current labels) + return tl.where( + valid_write, + union(labels_ptr, current_labels, neighbor_labels, valid_write), + current_labels, + ) + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_SIZE_H": 4, "BLOCK_SIZE_W": 16}, num_stages=1, num_warps=2 + ), + triton.Config( + {"BLOCK_SIZE_H": 4, "BLOCK_SIZE_W": 32}, num_stages=2, num_warps=4 + ), + ], + key=["H", "W"], + restore_value=["labels_ptr"], +) +@triton.jit +def _local_prop_kernel( + labels_ptr, + input_ptr, + H: tl.constexpr, + W: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_W: tl.constexpr, +): + # This is the meat of the Phase 2 to do local merging + # It will be launched with a 2D grid: + # - dim 0: batch index + # - dim 1: block index over HxW image (2D tiling) + pid_b = tl.program_id(0) + pid_hw = tl.program_id(1) + + # Calculate offsets for the core block + offsets_h = (pid_hw // tl.cdiv(W, BLOCK_SIZE_W)) * BLOCK_SIZE_H + tl.arange( + 0, BLOCK_SIZE_H + ) + offsets_w = (pid_hw % tl.cdiv(W, BLOCK_SIZE_W)) * BLOCK_SIZE_W + tl.arange( + 0, BLOCK_SIZE_W + ) + + base_offset = pid_b * H * W + offsets_2d = offsets_h[:, None] * W + offsets_w[None, :] + mask_2d = (offsets_h[:, None] < H) & (offsets_w[None, :] < W) + mask_1d = tl.ravel(mask_2d) + + # Load the current labels for the block - these are parent pointers + current_labels = tl.load( + labels_ptr + tl.ravel(base_offset + offsets_2d), mask=mask_1d, other=-1 + ) + current_values = tl.load( + input_ptr + base_offset + offsets_2d, mask=mask_2d, other=-1 + ) + valid_current = mask_1d & (current_labels != -1) + + # Horizontal merge + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + 0, + ) + # Vertical merge + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + 0, + -1, + ) + # Diagonal merges + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + -1, + ) + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + 1, + ) + + # This actually does some path compression, in a lightweight but beneficial way + tl.atomic_min( + labels_ptr + tl.ravel(base_offset + offsets_2d), current_labels, mask=mask_1d + ) + + +# ============================================================================== +# ## Phase 3: Pointer Jumping Kernel +# ============================================================================== +# This kernel performs pointer jumping to ensure that all pixels point directly to their root labels. +# This is done in a loop until convergence. + + +@triton.jit +def _pointer_jump_kernel( + labels_in_ptr, labels_out_ptr, numel: tl.constexpr, BLOCK_SIZE: tl.constexpr +): + """ + Pointer jumping kernel with double buffering to avoid race conditions. + Reads from labels_in_ptr and writes to labels_out_ptr. + """ + # This kernel is launched with a 1D grid, and does not care about batching explicitly. + # By construction, the labels are global indices across the batch, and we never perform + # cross-batch merges, so this is safe. + + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load current labels from input buffer + current_labels = tl.load(labels_in_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (current_labels != -1) + + # A mask to track which lanes have successfully completed their union. + done_mask = ~valid_mask + while tl_any(~(done_mask | ~valid_mask)): + parent_labels = tl.load( + labels_in_ptr + current_labels, mask=valid_mask, other=-1 + ) + + are_equal = current_labels == parent_labels + done_mask |= are_equal & valid_mask + + current_labels = tl.where( + ~done_mask, tl.minimum(current_labels, parent_labels), current_labels + ) + + # Write to output buffer (safe because we're not reading from it) + tl.store(labels_out_ptr + offsets, current_labels, mask=mask) + + +# ============================================================================== +# ## Phase 4: Kernels for Computing Component Sizes +# ============================================================================== + + +# Step 4.1: Count occurrences of each root label using atomic adds. +@triton.jit +def _count_labels_kernel(labels_ptr, sizes_ptr, numel, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load the final, converged labels + labels = tl.load(labels_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (labels != -1) + + # Atomically increment the counter for each label. This builds a histogram. + tl.atomic_add(sizes_ptr + labels, 1, mask=valid_mask) + + +# Step 4.2: Broadcast the computed sizes back to the output tensor. +@triton.jit +def _broadcast_sizes_kernel( + labels_ptr, sizes_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load the final labels + labels = tl.load(labels_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (labels != -1) + + # Look up the size for each label from the histogram + component_sizes = tl.load(sizes_ptr + labels, mask=valid_mask, other=0) + + # Write the size to the final output tensor. Background pixels get size 0. + tl.store(out_ptr + offsets, component_sizes, mask=mask) + + +def connected_components_triton(input_tensor: torch.Tensor): + """ + Computes connected components labeling on a batch of 2D integer tensors using Triton. + + Args: + input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted + + Returns: + Tuple[torch.Tensor, int]: A tuple containing: + - A BxHxW output tensor with dense labels. Background is 0. + - A BxHxW tensor with the size of the connected component for each pixel. + """ + assert ( + input_tensor.is_cuda and input_tensor.is_contiguous() + ), "Input tensor must be a contiguous CUDA tensor." + out_shape = input_tensor.shape + if input_tensor.dim() == 4 and input_tensor.shape[1] == 1: + input_tensor = input_tensor.squeeze(1) + else: + assert ( + input_tensor.dim() == 3 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + B, H, W = input_tensor.shape + numel = B * H * W + device = input_tensor.device + + # --- Allocate Tensors --- + labels = torch.empty_like(input_tensor, dtype=torch.int32) + output = torch.empty_like(input_tensor, dtype=torch.int32) + + # --- Phase 1 --- + BLOCK_SIZE = 256 + grid_init = (triton.cdiv(numel, BLOCK_SIZE),) + _init_labels_kernel[grid_init]( + input_tensor, + labels, + numel, + BLOCK_SIZE=BLOCK_SIZE, + ) + + # --- Phase 2 --- + grid_local_prop = lambda meta: ( + B, + triton.cdiv(H, meta["BLOCK_SIZE_H"]) * triton.cdiv(W, meta["BLOCK_SIZE_W"]), + ) + _local_prop_kernel[grid_local_prop](labels, input_tensor, H, W) + + # --- Phase 3 --- + BLOCK_SIZE = 256 + grid_jump = lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),) + _pointer_jump_kernel[grid_jump](labels, output, numel, BLOCK_SIZE=BLOCK_SIZE) + + # --- Phase 4 --- + # Allocate tensor to store the final output sizes + component_sizes_out = torch.empty_like(input_tensor, dtype=torch.int32) + + # Allocate a temporary 1D tensor to act as the histogram + # Size is numel because labels can be up to numel-1 + sizes_histogram = torch.zeros(numel, dtype=torch.int32, device=device) + + # 4.1: Count the occurrences of each label + grid_count = (triton.cdiv(numel, BLOCK_SIZE),) + _count_labels_kernel[grid_count]( + output, sizes_histogram, numel, BLOCK_SIZE=BLOCK_SIZE + ) + + # 2.2: Broadcast the counts to the final output tensor + grid_broadcast = (triton.cdiv(numel, BLOCK_SIZE),) + _broadcast_sizes_kernel[grid_broadcast]( + output, sizes_histogram, component_sizes_out, numel, BLOCK_SIZE=BLOCK_SIZE + ) + return output.view(out_shape) + 1, component_sizes_out.view(out_shape) diff --git a/sam3/perflib/triton/nms.py b/sam3/perflib/triton/nms.py new file mode 100644 index 0000000..ed800a1 --- /dev/null +++ b/sam3/perflib/triton/nms.py @@ -0,0 +1,124 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +# Adapted from https://github.com/stackav-oss/conch/blob/main/conch/kernels/vision/nms.py + +import torch +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({"cxpr_block_size": 128}), + triton.Config({"cxpr_block_size": 256}), + triton.Config({"cxpr_block_size": 512}), + triton.Config({"cxpr_block_size": 1024}), + triton.Config({"cxpr_block_size": 2048}), + triton.Config({"cxpr_block_size": 4096}), + triton.Config({"cxpr_block_size": 8192}), + ], + key=["num_boxes"], +) +@triton.jit +def _nms_suppression_kernel( + # Tensors + iou_mask_ptr: tl.tensor, # [N, N] + keep_mask_ptr: tl.tensor, # [N] + # Scalars + num_boxes: tl.int32, + # Strides + iou_mask_stride: tl.int32, + # Constexprs + cxpr_block_size: tl.constexpr, +) -> None: + """NMS suppression kernel. + + Args: + iou_mask_ptr: Pointer to precomputed IoU mask, shape: (N, N). + keep_mask_ptr: Pointer to keep mask tensor, shape: (N,). + num_boxes: Number of boxes. + iou_mask_stride: Stride for IoU mask tensor. + cxpr_block_size: Block size for processing. + """ + # Sequential NMS: for each box in sorted order, suppress later boxes + for current_box_idx in range(num_boxes - 1): + # Check if current box is still kept + is_kept = tl.load(keep_mask_ptr + current_box_idx) + if is_kept: + # IoU mask row offset for the current box + # Because the IoU mask is sorted by score, we will only consider boxes that come after the current box. + # This means we only need to read the upper triangular part of the IoU mask. + iou_row_offset = current_box_idx * iou_mask_stride + + # Only process boxes that come after the current box + next_box_idx = current_box_idx + 1 + remaining_boxes = num_boxes - next_box_idx + + # Iterate blockwise through the columns + for block_idx in range(tl.cdiv(remaining_boxes, cxpr_block_size)): + # Masked load of indices for the target boxes in the current block + block_start = next_box_idx + block_idx * cxpr_block_size + target_box_offsets = block_start + tl.arange(0, cxpr_block_size) + target_box_mask = target_box_offsets < num_boxes + + # Suppress boxes with lower scores that have high IoU + suppression_mask = tl.load( + iou_mask_ptr + iou_row_offset + target_box_offsets, + mask=target_box_mask, + other=False, + ) + suppression_mask = tl.cast(suppression_mask, tl.int1) + + # Conditionally store suppression result for high-IoU boxes + tl.store( + keep_mask_ptr + target_box_offsets, False, mask=suppression_mask + ) + + # Potential race condition: we need to ensure all threads complete the store before the next + # iteration otherwise we may load stale data for whether or not a box has been suppressed. + tl.debug_barrier() + + +def nms_triton( + ious: torch.Tensor, + scores: torch.Tensor, + iou_threshold: float, +) -> torch.Tensor: + """Perform NMS given the iou matrix, the scores and the iou threshold + + Args: + ious: Pairwise IoU tensor of shape (N, N). + scores: Scores tensor of shape (N,). + iou_threshold: IoU threshold for suppression. + + Returns: + Tensor: Indices of kept boxes, sorted by decreasing score. + """ + assert scores.dim() == 1, "Scores must be 1D" + iou_mask = ious > iou_threshold + assert iou_mask.dim() == 2 + assert iou_mask.shape[0] == iou_mask.shape[1] == scores.shape[0] + assert iou_mask.device == scores.device + assert iou_mask.dtype == torch.bool + + num_boxes = scores.size(0) + keep_mask = torch.ones(len(scores), device=scores.device, dtype=torch.bool) + + # Sort boxes by scores in descending order + _, sorted_indices = torch.sort(scores, dim=0, stable=True, descending=True) + iou_mask = iou_mask[sorted_indices][:, sorted_indices].contiguous() + + # For the suppression stage, we need to process sequentially, but we'll still take + # advantage of parallelism by processing in blocks in one program. + stage2_grid = (1,) + _nms_suppression_kernel[stage2_grid]( + # Tensors + iou_mask_ptr=iou_mask, + keep_mask_ptr=keep_mask, + # Scalars + num_boxes=num_boxes, + # Strides + iou_mask_stride=iou_mask.stride(0), + ) + # Extract indices of kept boxes + return sorted_indices[keep_mask] diff --git a/sam3/sam/__init__.py b/sam3/sam/__init__.py new file mode 100644 index 0000000..8b35da6 --- /dev/null +++ b/sam3/sam/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +from .mask_decoder import MaskDecoder +from .prompt_encoder import PromptEncoder +from .transformer import TwoWayTransformer diff --git a/sam3/sam/common.py b/sam3/sam/common.py new file mode 100644 index 0000000..b6d1587 --- /dev/null +++ b/sam3/sam/common.py @@ -0,0 +1,39 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import Type + +import torch +import torch.nn as nn + + +class MLPBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + mlp_dim: int, + act: Type[nn.Module] = nn.GELU, + ) -> None: + super().__init__() + self.lin1 = nn.Linear(embedding_dim, mlp_dim) + self.lin2 = nn.Linear(mlp_dim, embedding_dim) + self.act = act() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lin2(self.act(self.lin1(x))) + + +# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa +# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa +class LayerNorm2d(nn.Module): + def __init__(self, num_channels: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(num_channels)) + self.bias = nn.Parameter(torch.zeros(num_channels)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x diff --git a/sam3/sam/mask_decoder.py b/sam3/sam/mask_decoder.py new file mode 100644 index 0000000..b4ac397 --- /dev/null +++ b/sam3/sam/mask_decoder.py @@ -0,0 +1,319 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import List, Optional, Tuple, Type + +import torch +from torch import nn +from torch.nn import functional as F + +from .common import LayerNorm2d + + +class MaskDecoder(nn.Module): + def __init__( + self, + *, + transformer_dim: int, + transformer: nn.Module, + num_multimask_outputs: int = 3, + activation: Type[nn.Module] = nn.GELU, + iou_head_depth: int = 3, + iou_head_hidden_dim: int = 256, + use_high_res_features: bool = False, + iou_prediction_use_sigmoid=False, + dynamic_multimask_via_stability=False, + dynamic_multimask_stability_delta=0.05, + dynamic_multimask_stability_thresh=0.98, + pred_obj_scores: bool = False, + pred_obj_scores_mlp: bool = False, + use_multimask_token_for_obj_ptr: bool = False, + ) -> None: + """ + Predicts masks given an image and prompt embeddings, using a + transformer architecture. + + Arguments: + transformer_dim (int): the channel dimension of the transformer + transformer (nn.Module): the transformer used to predict masks + num_multimask_outputs (int): the number of masks to predict + when disambiguating masks + activation (nn.Module): the type of activation to use when + upscaling masks + iou_head_depth (int): the depth of the MLP used to predict + mask quality + iou_head_hidden_dim (int): the hidden dimension of the MLP + used to predict mask quality + """ + super().__init__() + self.transformer_dim = transformer_dim + self.transformer = transformer + + self.num_multimask_outputs = num_multimask_outputs + + self.iou_token = nn.Embedding(1, transformer_dim) + self.num_mask_tokens = num_multimask_outputs + 1 + self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) + + self.pred_obj_scores = pred_obj_scores + if self.pred_obj_scores: + self.obj_score_token = nn.Embedding(1, transformer_dim) + self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr + + self.output_upscaling = nn.Sequential( + nn.ConvTranspose2d( + transformer_dim, transformer_dim // 4, kernel_size=2, stride=2 + ), + LayerNorm2d(transformer_dim // 4), + activation(), + nn.ConvTranspose2d( + transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2 + ), + activation(), + ) + self.use_high_res_features = use_high_res_features + if use_high_res_features: + self.conv_s0 = nn.Conv2d( + transformer_dim, transformer_dim // 8, kernel_size=1, stride=1 + ) + self.conv_s1 = nn.Conv2d( + transformer_dim, transformer_dim // 4, kernel_size=1, stride=1 + ) + + self.output_hypernetworks_mlps = nn.ModuleList( + [ + MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) + for i in range(self.num_mask_tokens) + ] + ) + + self.iou_prediction_head = MLP( + transformer_dim, + iou_head_hidden_dim, + self.num_mask_tokens, + iou_head_depth, + sigmoid_output=iou_prediction_use_sigmoid, + ) + if self.pred_obj_scores: + self.pred_obj_score_head = nn.Linear(transformer_dim, 1) + if pred_obj_scores_mlp: + self.pred_obj_score_head = MLP(transformer_dim, transformer_dim, 1, 3) + + # When outputting a single mask, optionally we can dynamically fall back to the best + # multimask output token if the single mask output token gives low stability scores. + self.dynamic_multimask_via_stability = dynamic_multimask_via_stability + self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta + self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh + + def forward( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + multimask_output: bool, + repeat_image: bool, + high_res_features: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Arguments: + image_embeddings (torch.Tensor): the embeddings from the image encoder + image_pe (torch.Tensor): positional encoding with the shape of image_embeddings + sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes + dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs + multimask_output (bool): Whether to return multiple masks or a single + mask. + + Returns: + torch.Tensor: batched predicted masks + torch.Tensor: batched predictions of mask quality + torch.Tensor: batched SAM token for mask output + """ + masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks( + image_embeddings=image_embeddings, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_prompt_embeddings, + dense_prompt_embeddings=dense_prompt_embeddings, + repeat_image=repeat_image, + high_res_features=high_res_features, + ) + + # Select the correct mask or masks for output + if multimask_output: + masks = masks[:, 1:, :, :] + iou_pred = iou_pred[:, 1:] + elif self.dynamic_multimask_via_stability and not self.training: + masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred) + else: + masks = masks[:, 0:1, :, :] + iou_pred = iou_pred[:, 0:1] + + if multimask_output and self.use_multimask_token_for_obj_ptr: + sam_tokens_out = mask_tokens_out[:, 1:] # [b, 3, c] shape + else: + # Take the mask output token. Here we *always* use the token for single mask output. + # At test time, even if we track after 1-click (and using multimask_output=True), + # we still take the single mask token here. The rationale is that we always track + # after multiple clicks during training, so the past tokens seen during training + # are always the single mask token (and we'll let it be the object-memory token). + sam_tokens_out = mask_tokens_out[:, 0:1] # [b, 1, c] shape + + # Prepare output + return masks, iou_pred, sam_tokens_out, object_score_logits + + def predict_masks( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + repeat_image: bool, + high_res_features: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predicts masks. See 'forward' for more details.""" + # Concatenate output tokens + s = 0 + if self.pred_obj_scores: + output_tokens = torch.cat( + [ + self.obj_score_token.weight, + self.iou_token.weight, + self.mask_tokens.weight, + ], + dim=0, + ) + s = 1 + else: + output_tokens = torch.cat( + [self.iou_token.weight, self.mask_tokens.weight], dim=0 + ) + output_tokens = output_tokens.unsqueeze(0).expand( + sparse_prompt_embeddings.size(0), -1, -1 + ) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + + # Expand per-image data in batch direction to be per-mask + if repeat_image: + src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0) + else: + assert image_embeddings.shape[0] == tokens.shape[0] + src = image_embeddings + src = src + dense_prompt_embeddings + assert ( + image_pe.size(0) == 1 + ), "image_pe should have size 1 in batch dim (from `get_dense_pe()`)" + pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0) + b, c, h, w = src.shape + + # Run the transformer + hs, src = self.transformer(src, pos_src, tokens) + iou_token_out = hs[:, s, :] + mask_tokens_out = hs[:, s + 1 : (s + 1 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + src = src.transpose(1, 2).view(b, c, h, w) + if not self.use_high_res_features: + upscaled_embedding = self.output_upscaling(src) + else: + dc1, ln1, act1, dc2, act2 = self.output_upscaling + feat_s0, feat_s1 = high_res_features + upscaled_embedding = act1(ln1(dc1(src) + feat_s1)) + upscaled_embedding = act2(dc2(upscaled_embedding) + feat_s0) + + hyper_in_list: List[torch.Tensor] = [] + for i in range(self.num_mask_tokens): + hyper_in_list.append( + self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]) + ) + hyper_in = torch.stack(hyper_in_list, dim=1) + b, c, h, w = upscaled_embedding.shape + masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) + + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + if self.pred_obj_scores: + assert s == 1 + object_score_logits = self.pred_obj_score_head(hs[:, 0, :]) + else: + # Obj scores logits - default to 10.0, i.e. assuming the object is present, sigmoid(10)=1 + object_score_logits = 10.0 * iou_pred.new_ones(iou_pred.shape[0], 1) + + return masks, iou_pred, mask_tokens_out, object_score_logits + + def _get_stability_scores(self, mask_logits): + """ + Compute stability scores of the mask logits based on the IoU between upper and + lower thresholds. + """ + mask_logits = mask_logits.flatten(-2) + stability_delta = self.dynamic_multimask_stability_delta + area_i = torch.sum(mask_logits > stability_delta, dim=-1).float() + area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float() + stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0) + return stability_scores + + def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores): + """ + When outputting a single mask, if the stability score from the current single-mask + output (based on output token 0) falls below a threshold, we instead select from + multi-mask outputs (based on output token 1~3) the mask with the highest predicted + IoU score. This is intended to ensure a valid mask for both clicking and tracking. + """ + # The best mask from multimask output tokens (1~3) + multimask_logits = all_mask_logits[:, 1:, :, :] + multimask_iou_scores = all_iou_scores[:, 1:] + best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1) + batch_inds = torch.arange( + multimask_iou_scores.size(0), device=all_iou_scores.device + ) + best_multimask_logits = multimask_logits[batch_inds, best_scores_inds] + best_multimask_logits = best_multimask_logits.unsqueeze(1) + best_multimask_iou_scores = multimask_iou_scores[batch_inds, best_scores_inds] + best_multimask_iou_scores = best_multimask_iou_scores.unsqueeze(1) + + # The mask from singlemask output token 0 and its stability score + singlemask_logits = all_mask_logits[:, 0:1, :, :] + singlemask_iou_scores = all_iou_scores[:, 0:1] + stability_scores = self._get_stability_scores(singlemask_logits) + is_stable = stability_scores >= self.dynamic_multimask_stability_thresh + + # Dynamically fall back to best multimask output upon low stability scores. + mask_logits_out = torch.where( + is_stable[..., None, None].expand_as(singlemask_logits), + singlemask_logits, + best_multimask_logits, + ) + iou_scores_out = torch.where( + is_stable.expand_as(singlemask_iou_scores), + singlemask_iou_scores, + best_multimask_iou_scores, + ) + return mask_logits_out, iou_scores_out + + +# Lightly adapted from +# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa +class MLP(nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + sigmoid_output: bool = False, + ) -> None: + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + self.sigmoid_output = sigmoid_output + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + if self.sigmoid_output: + x = F.sigmoid(x) + return x diff --git a/sam3/sam/prompt_encoder.py b/sam3/sam/prompt_encoder.py new file mode 100644 index 0000000..145ea9f --- /dev/null +++ b/sam3/sam/prompt_encoder.py @@ -0,0 +1,243 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from typing import Any, Optional, Tuple, Type + +import numpy as np +import torch +from torch import nn + +from .common import LayerNorm2d + + +class PromptEncoder(nn.Module): + def __init__( + self, + embed_dim: int, + image_embedding_size: Tuple[int, int], + input_image_size: Tuple[int, int], + mask_in_chans: int, + activation: Type[nn.Module] = nn.GELU, + ) -> None: + """ + Encodes prompts for input to SAM's mask decoder. + + Arguments: + embed_dim (int): The prompts' embedding dimension + image_embedding_size (tuple(int, int)): The spatial size of the + image embedding, as (H, W). + input_image_size (int): The padded size of the image as input + to the image encoder, as (H, W). + mask_in_chans (int): The number of hidden channels used for + encoding input masks. + activation (nn.Module): The activation to use when encoding + input masks. + """ + super().__init__() + self.embed_dim = embed_dim + self.input_image_size = input_image_size + self.image_embedding_size = image_embedding_size + self.pe_layer = PositionEmbeddingRandom(embed_dim // 2) + + self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners + point_embeddings = [ + nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings) + ] + self.point_embeddings = nn.ModuleList(point_embeddings) + self.not_a_point_embed = nn.Embedding(1, embed_dim) + + self.mask_input_size = ( + 4 * image_embedding_size[0], + 4 * image_embedding_size[1], + ) + self.mask_downscaling = nn.Sequential( + nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans // 4), + activation(), + nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans), + activation(), + nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1), + ) + self.no_mask_embed = nn.Embedding(1, embed_dim) + + def get_dense_pe(self) -> torch.Tensor: + """ + Returns the positional encoding used to encode point prompts, + applied to a dense set of points the shape of the image encoding. + + Returns: + torch.Tensor: Positional encoding with shape + 1x(embed_dim)x(embedding_h)x(embedding_w) + """ + return self.pe_layer(self.image_embedding_size).unsqueeze(0) + + def _embed_points( + self, + points: torch.Tensor, + labels: torch.Tensor, + pad: bool, + ) -> torch.Tensor: + """Embeds point prompts.""" + points = points + 0.5 # Shift to center of pixel + if pad: + padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device) + padding_label = -torch.ones((labels.shape[0], 1), device=labels.device) + points = torch.cat([points, padding_point], dim=1) + labels = torch.cat([labels, padding_label], dim=1) + point_embedding = self.pe_layer.forward_with_coords( + points, self.input_image_size + ) + + point_embedding = torch.where( + (labels == -1).unsqueeze(-1), + torch.zeros_like(point_embedding) + self.not_a_point_embed.weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 0).unsqueeze(-1), + point_embedding + self.point_embeddings[0].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 1).unsqueeze(-1), + point_embedding + self.point_embeddings[1].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 2).unsqueeze(-1), + point_embedding + self.point_embeddings[2].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 3).unsqueeze(-1), + point_embedding + self.point_embeddings[3].weight, + point_embedding, + ) + return point_embedding + + def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor: + """Embeds box prompts.""" + boxes = boxes + 0.5 # Shift to center of pixel + coords = boxes.reshape(-1, 2, 2) + corner_embedding = self.pe_layer.forward_with_coords( + coords, self.input_image_size + ) + corner_embedding[:, 0, :] += self.point_embeddings[2].weight + corner_embedding[:, 1, :] += self.point_embeddings[3].weight + return corner_embedding + + def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor: + """Embeds mask inputs.""" + mask_embedding = self.mask_downscaling(masks) + return mask_embedding + + def _get_batch_size( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> int: + """ + Gets the batch size of the output given the batch size of the input prompts. + """ + if points is not None: + return points[0].shape[0] + elif boxes is not None: + return boxes.shape[0] + elif masks is not None: + return masks.shape[0] + else: + return 1 + + def _get_device(self) -> torch.device: + return self.point_embeddings[0].weight.device + + def forward( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Embeds different types of prompts, returning both sparse and dense + embeddings. + + Arguments: + points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates + and labels to embed. + boxes (torch.Tensor or none): boxes to embed + masks (torch.Tensor or none): masks to embed + + Returns: + torch.Tensor: sparse embeddings for the points and boxes, with shape + BxNx(embed_dim), where N is determined by the number of input points + and boxes. + torch.Tensor: dense embeddings for the masks, in the shape + Bx(embed_dim)x(embed_H)x(embed_W) + """ + bs = self._get_batch_size(points, boxes, masks) + sparse_embeddings = torch.empty( + (bs, 0, self.embed_dim), device=self._get_device() + ) + if points is not None: + coords, labels = points + point_embeddings = self._embed_points(coords, labels, pad=(boxes is None)) + sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1) + if boxes is not None: + box_embeddings = self._embed_boxes(boxes) + sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1) + + if masks is not None: + dense_embeddings = self._embed_masks(masks) + else: + dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand( + bs, -1, self.image_embedding_size[0], self.image_embedding_size[1] + ) + + return sparse_embeddings, dense_embeddings + + +class PositionEmbeddingRandom(nn.Module): + """ + Positional encoding using random spatial frequencies. + """ + + def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: + super().__init__() + if scale is None or scale <= 0.0: + scale = 1.0 + self.register_buffer( + "positional_encoding_gaussian_matrix", + scale * torch.randn((2, num_pos_feats)), + ) + + def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: + """Positionally encode points that are normalized to [0,1].""" + # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape + coords = 2 * coords - 1 + coords = coords @ self.positional_encoding_gaussian_matrix + coords = 2 * np.pi * coords + # outputs d_1 x ... x d_n x C shape + return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) + + def forward(self, size: Tuple[int, int]) -> torch.Tensor: + """Generate positional encoding for a grid of the specified size.""" + h, w = size + device: Any = self.positional_encoding_gaussian_matrix.device + grid = torch.ones((h, w), device=device, dtype=torch.float32) + y_embed = grid.cumsum(dim=0) - 0.5 + x_embed = grid.cumsum(dim=1) - 0.5 + y_embed = y_embed / h + x_embed = x_embed / w + + pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) + return pe.permute(2, 0, 1) # C x H x W + + def forward_with_coords( + self, coords_input: torch.Tensor, image_size: Tuple[int, int] + ) -> torch.Tensor: + """Positionally encode points that are not normalized to [0,1].""" + coords = coords_input.clone() + coords[:, :, 0] = coords[:, :, 0] / image_size[1] + coords[:, :, 1] = coords[:, :, 1] / image_size[0] + return self._pe_encoding(coords.to(torch.float)) # B x N x C diff --git a/sam3/sam/rope.py b/sam3/sam/rope.py new file mode 100644 index 0000000..2db01b6 --- /dev/null +++ b/sam3/sam/rope.py @@ -0,0 +1,161 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +""" +Adapted from: +1. https://github.com/meta-llama/codellama/blob/main/llama/model.py +2. https://github.com/naver-ai/rope-vit +3. https://github.com/lucidrains/rotary-embedding-torch +""" + +from typing import Optional + +import torch +from einops import rearrange, repeat +from torch import broadcast_tensors, nn + + +def init_t_xy(end_x: int, end_y: int, scale: float = 1.0, offset: int = 0, device=None): + t = torch.arange(end_x * end_y, dtype=torch.float32, device=device) + t_x = (t % end_x).float() + t_y = torch.div(t, end_x, rounding_mode="floor").float() + return t_x * scale + offset, t_y * scale + offset + + +def compute_axial_cis( + dim: int, + end_x: int, + end_y: int, + theta: float = 10000.0, + scale_pos: float = 1.0, + offset: int = 0, + device=None, +): + freqs_x = 1.0 / ( + theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim) + ) + freqs_y = 1.0 / ( + theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim) + ) + + t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset, device=device) + freqs_x = torch.outer(t_x, freqs_x) + freqs_y = torch.outer(t_y, freqs_y) + freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x) + freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y) + return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1) + + +def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): + ndim = x.ndim + assert 0 <= 1 < ndim + assert freqs_cis.shape == (x.shape[-2], x.shape[-1]) + shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def apply_rotary_enc( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: torch.Tensor, + repeat_freqs_k: bool = False, +): + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) + xk_ = ( + torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) + if xk.shape[-2] != 0 + else None + ) + freqs_cis = reshape_for_broadcast(freqs_cis, xq_) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) + if xk_ is None: + # no keys to rotate, due to dropout + return xq_out.type_as(xq).to(xq.device), xk + # repeat freqs along seq_len dim to match k seq_len + if repeat_freqs_k: + r = xk_.shape[-2] // xq_.shape[-2] + freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1) + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) + return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device) + + +def complex_mult(xq_real, xq_imag, freqs_cis_real, freqs_cis_imag): + # Compute the real part of the product + real_part = xq_real * freqs_cis_real - xq_imag * freqs_cis_imag + # Compute the imaginary part of the product + imag_part = xq_real * freqs_cis_imag + xq_imag * freqs_cis_real + # Stack the real and imaginary parts along the last dimension + return torch.stack([real_part, imag_part], dim=-1) + + +def apply_rotary_enc_real( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis_real: torch.Tensor, + freqs_cis_imag: torch.Tensor, + repeat_freqs_k: bool = False, +): + assert xk is not None + assert xk.shape[-2] != 0 + + xq_real = xq.float().reshape(*xq.shape[:-1], -1, 2)[..., 0] + xq_imag = xq.float().reshape(*xq.shape[:-1], -1, 2)[..., 1] + xk_real = xk.float().reshape(*xk.shape[:-1], -1, 2)[..., 0] + xk_imag = xk.float().reshape(*xk.shape[:-1], -1, 2)[..., 1] + freqs_cis_real = reshape_for_broadcast(freqs_cis_real, xq_real) + freqs_cis_imag = reshape_for_broadcast(freqs_cis_imag, xq_imag) + xq_out = complex_mult(xq_real, xq_imag, freqs_cis_real, freqs_cis_imag).flatten(3) + if repeat_freqs_k: + r = xk_real.shape[-2] // xq_real.shape[-2] + freqs_cis_real = freqs_cis_real.repeat(*([1] * (freqs_cis_real.ndim - 2)), r, 1) + freqs_cis_imag = freqs_cis_imag.repeat(*([1] * (freqs_cis_imag.ndim - 2)), r, 1) + xk_out = complex_mult(xk_real, xk_imag, freqs_cis_real, freqs_cis_imag).flatten(3) + # xq_out = torch.view_as_real(torch.complex(xq_real, xq_imag) * torch.complex(freqs_cis_real, freqs_cis_imag)).flatten(3) + # xk_out = torch.view_as_real(torch.compelx(xk_real, xk_imag) * torch.complex(freqs_cis_real, freqs_cis_imag)).flatten(3) + return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device) + + +# rotary embedding helper functions +def broadcat(tensors, dim=-1): + broadcasted_tensors = broadcast_tensors(*tensors) + return torch.cat(broadcasted_tensors, dim=dim) + + +def rotate_half(x: torch.Tensor): + x = rearrange(x, "... (d r) -> ... d r", r=2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return rearrange(x, "... d r -> ... (d r)") + + +class VisionRotaryEmbeddingVE(nn.Module): + def __init__( + self, + dim: int, + seq_len: int, + pt_seq_len: Optional[int] = None, + theta: float = 10000.0, + offset: int = 1, # specific to VE + ): + super().__init__() + + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + scale = 1.0 + if pt_seq_len is not None: + scale = pt_seq_len / seq_len + + # offset of +1 following VE - even though for the + # attention op only differences matter + t = torch.arange(seq_len) * scale + offset + + freqs = torch.einsum("..., f -> ... f", t, freqs) + freqs = repeat(freqs, "... n -> ... (n r)", r=2) + + freqs = broadcat((freqs[None, :, :], freqs[:, None, :]), dim=-1) + freqs_cos = freqs.cos().view(-1, freqs.shape[-1]) + freqs_sin = freqs.sin().view(-1, freqs.shape[-1]) + + self.register_buffer("freqs_cos", freqs_cos) + self.register_buffer("freqs_sin", freqs_sin) + + def forward(self, t: torch.Tensor): + return t * self.freqs_cos + rotate_half(t) * self.freqs_sin diff --git a/sam3/sam/transformer.py b/sam3/sam/transformer.py new file mode 100644 index 0000000..3e96c28 --- /dev/null +++ b/sam3/sam/transformer.py @@ -0,0 +1,358 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import math +from functools import partial +from typing import Tuple, Type + +import torch +import torch.nn.functional as F + +from sam3.sam.rope import apply_rotary_enc, apply_rotary_enc_real, compute_axial_cis +from torch import nn, Tensor + +from .common import MLPBlock + + +class TwoWayTransformer(nn.Module): + def __init__( + self, + depth: int, + embedding_dim: int, + num_heads: int, + mlp_dim: int, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + ) -> None: + """ + A transformer decoder that attends to an input image using + queries whose positional embedding is supplied. + + Args: + depth (int): number of layers in the transformer + embedding_dim (int): the channel dimension for the input embeddings + num_heads (int): the number of heads for multihead attention. Must + divide embedding_dim + mlp_dim (int): the channel dimension internal to the MLP block + activation (nn.Module): the activation to use in the MLP block + """ + super().__init__() + self.depth = depth + self.embedding_dim = embedding_dim + self.num_heads = num_heads + self.mlp_dim = mlp_dim + self.layers = nn.ModuleList() + + for i in range(depth): + self.layers.append( + TwoWayAttentionBlock( + embedding_dim=embedding_dim, + num_heads=num_heads, + mlp_dim=mlp_dim, + activation=activation, + attention_downsample_rate=attention_downsample_rate, + skip_first_layer_pe=(i == 0), + ) + ) + + self.final_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm_final_attn = nn.LayerNorm(embedding_dim) + + def forward( + self, + image_embedding: Tensor, + image_pe: Tensor, + point_embedding: Tensor, + ) -> Tuple[Tensor, Tensor]: + """ + Args: + image_embedding (torch.Tensor): image to attend to. Should be shape + B x embedding_dim x h x w for any h and w. + image_pe (torch.Tensor): the positional encoding to add to the image. Must + have the same shape as image_embedding. + point_embedding (torch.Tensor): the embedding to add to the query points. + Must have shape B x N_points x embedding_dim for any N_points. + + Returns: + torch.Tensor: the processed point_embedding + torch.Tensor: the processed image_embedding + """ + # BxCxHxW -> BxHWxC == B x N_image_tokens x C + bs, c, h, w = image_embedding.shape + image_embedding = image_embedding.flatten(2).permute(0, 2, 1) + image_pe = image_pe.flatten(2).permute(0, 2, 1) + + # Prepare queries + queries = point_embedding + keys = image_embedding + + # Apply transformer blocks and final layernorm + for layer in self.layers: + queries, keys = layer( + queries=queries, + keys=keys, + query_pe=point_embedding, + key_pe=image_pe, + ) + + # Apply the final attention layer from the points to the image + q = queries + point_embedding + k = keys + image_pe + attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm_final_attn(queries) + + return queries, keys + + +class TwoWayAttentionBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + num_heads: int, + mlp_dim: int = 2048, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + skip_first_layer_pe: bool = False, + ) -> None: + """ + A transformer block with four layers: (1) self-attention of sparse + inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp + block on sparse inputs, and (4) cross attention of dense inputs to sparse + inputs. + + Arguments: + embedding_dim (int): the channel dimension of the embeddings + num_heads (int): the number of heads in the attention layers + mlp_dim (int): the hidden dimension of the mlp block + activation (nn.Module): the activation of the mlp block + skip_first_layer_pe (bool): skip the PE on the first layer + """ + super().__init__() + self.self_attn = Attention(embedding_dim, num_heads) + self.norm1 = nn.LayerNorm(embedding_dim) + + self.cross_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm2 = nn.LayerNorm(embedding_dim) + + self.mlp = MLPBlock(embedding_dim, mlp_dim, activation) + self.norm3 = nn.LayerNorm(embedding_dim) + + self.norm4 = nn.LayerNorm(embedding_dim) + self.cross_attn_image_to_token = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + + self.skip_first_layer_pe = skip_first_layer_pe + + def forward( + self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor + ) -> Tuple[Tensor, Tensor]: + # Self attention block + if self.skip_first_layer_pe: + queries = self.self_attn(q=queries, k=queries, v=queries) + else: + q = queries + query_pe + attn_out = self.self_attn(q=q, k=q, v=queries) + queries = queries + attn_out + queries = self.norm1(queries) + + # Cross attention block, tokens attending to image embedding + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm2(queries) + + # MLP block + mlp_out = self.mlp(queries) + queries = queries + mlp_out + queries = self.norm3(queries) + + # Cross attention block, image embedding attending to tokens + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries) + keys = keys + attn_out + keys = self.norm4(keys) + + return queries, keys + + +class Attention(nn.Module): + """ + An attention layer that allows for downscaling the size of the embedding + after projection to queries, keys, and values. + """ + + def __init__( + self, + embedding_dim: int, + num_heads: int, + downsample_rate: int = 1, + dropout: float = 0.0, + kv_in_dim: int = None, + use_fa3: bool = False, + ) -> None: + super().__init__() + self.embedding_dim = embedding_dim + self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim + self.internal_dim = embedding_dim // downsample_rate + self.num_heads = num_heads + self.use_fa3 = use_fa3 + assert ( + self.internal_dim % num_heads == 0 + ), "num_heads must divide embedding_dim." + + self.q_proj = nn.Linear(embedding_dim, self.internal_dim) + self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim) + self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim) + self.out_proj = nn.Linear(self.internal_dim, embedding_dim) + + self.dropout_p = dropout + + def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor: + b, n, c = x.shape + x = x.reshape(b, n, num_heads, c // num_heads) + return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head + + def _recombine_heads(self, x: Tensor) -> Tensor: + b, n_heads, n_tokens, c_per_head = x.shape + x = x.transpose(1, 2) + return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: + # Input projections + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # Separate into heads + q = self._separate_heads(q, self.num_heads) + k = self._separate_heads(k, self.num_heads) + v = self._separate_heads(v, self.num_heads) + + dropout_p = self.dropout_p if self.training else 0.0 + # Attention + # with torch.backends.cuda.sdp_kernel( + # enable_flash=USE_FLASH_ATTN, + # # if Flash attention kernel is off, then math kernel needs to be enabled + # enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON, + # enable_mem_efficient=OLD_GPU, + # ): + # Let's trust the dispatcher.... + if self.use_fa3: + from sam3.perflib.fa3 import flash_attn_func + + assert dropout_p == 0.0 + out = flash_attn_func( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ).transpose(1, 2) + else: + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(True) + out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p) + + out = self._recombine_heads(out) + out = self.out_proj(out) + + return out + + +class RoPEAttention(Attention): + """Attention with rotary position encoding.""" + + def __init__( + self, + *args, + rope_theta=10000.0, + # whether to repeat q rope to match k length + # this is needed for cross-attention to memories + rope_k_repeat=False, + feat_sizes=(64, 64), # [w, h] for stride 16 feats at 1024 resolution + use_rope_real=False, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.use_rope_real = use_rope_real + self.compute_cis = partial( + compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta + ) + device = torch.device("cuda") if torch.cuda.is_available() else None + self.freqs_cis = self.compute_cis( + end_x=feat_sizes[0], end_y=feat_sizes[1], device=device + ) + if self.use_rope_real: + self.freqs_cis_real = self.freqs_cis.real + self.freqs_cis_imag = self.freqs_cis.imag + self.rope_k_repeat = rope_k_repeat + + def forward( + self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0 + ) -> Tensor: + # Input projections + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # Separate into heads + q = self._separate_heads(q, self.num_heads) + k = self._separate_heads(k, self.num_heads) + v = self._separate_heads(v, self.num_heads) + + # Apply rotary position encoding + w = h = math.sqrt(q.shape[-2]) + if self.freqs_cis.shape[0] != q.shape[-2]: + self.freqs_cis = self.compute_cis(end_x=w, end_y=h, device=q.device) + self.freqs_cis_real = self.freqs_cis.real + self.freqs_cis_imag = self.freqs_cis.imag + if q.shape[-2] != k.shape[-2]: + assert self.rope_k_repeat + + num_k_rope = k.size(-2) - num_k_exclude_rope + if self.use_rope_real: + q, k[:, :, :num_k_rope] = apply_rotary_enc_real( + q, + k[:, :, :num_k_rope], + freqs_cis_real=self.freqs_cis_real, + freqs_cis_imag=self.freqs_cis_imag, + repeat_freqs_k=self.rope_k_repeat, + ) + else: + q, k[:, :, :num_k_rope] = apply_rotary_enc( + q, + k[:, :, :num_k_rope], + self.freqs_cis, + repeat_freqs_k=self.rope_k_repeat, + ) + + dropout_p = self.dropout_p if self.training else 0.0 + # Attention + # with torch.backends.cuda.sdp_kernel( + # enable_flash=USE_FLASH_ATTN, + # # if Flash attention kernel is off, then math kernel needs to be enabled + # enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON, + # enable_mem_efficient=OLD_GPU, + # ): + # Let's trust the dispatcher.... + if self.use_fa3: + from sam3.perflib.fa3 import flash_attn_func + + assert dropout_p == 0.0 + out = flash_attn_func( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ).transpose(1, 2) + else: + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(True) + out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p) + + out = self._recombine_heads(out) + out = self.out_proj(out) + + return out