(
mask_weights: Dict[str, torch.Tensor],
prompt_weights: Dict[str, torch.Tensor],
image_embeddings: torch.Tensor,
neck_trk_0: torch.Tensor,
neck_trk_1: torch.Tensor,
)
| 304 | |
| 305 | |
| 306 | def run_sam_decoder( |
| 307 | mask_weights: Dict[str, torch.Tensor], |
| 308 | prompt_weights: Dict[str, torch.Tensor], |
| 309 | image_embeddings: torch.Tensor, |
| 310 | neck_trk_0: torch.Tensor, |
| 311 | neck_trk_1: torch.Tensor, |
| 312 | ) -> dict[str, torch.Tensor]: |
| 313 | sparse_embeddings, dense_embeddings, image_pe = build_no_point_prompt_encoder_outputs(prompt_weights) |
| 314 | |
| 315 | feat_s0 = F.conv2d( |
| 316 | neck_trk_0, |
| 317 | mask_weights["conv_s0.weight"].float(), |
| 318 | mask_weights["conv_s0.bias"].float(), |
| 319 | ) |
| 320 | feat_s1 = F.conv2d( |
| 321 | neck_trk_1, |
| 322 | mask_weights["conv_s1.weight"].float(), |
| 323 | mask_weights["conv_s1.bias"].float(), |
| 324 | ) |
| 325 | |
| 326 | output_tokens = torch.cat( |
| 327 | [ |
| 328 | mask_weights["obj_score_token.weight"].float(), |
| 329 | mask_weights["iou_token.weight"].float(), |
| 330 | mask_weights["mask_tokens.weight"].float(), |
| 331 | ], |
| 332 | dim=0, |
| 333 | ).unsqueeze(0) |
| 334 | tokens = torch.cat((output_tokens, sparse_embeddings), dim=1) |
| 335 | |
| 336 | src = image_embeddings + dense_embeddings |
| 337 | pos_src = image_pe |
| 338 | bsz, dim, h, w = src.shape |
| 339 | keys = src.flatten(2).permute(0, 2, 1) |
| 340 | key_pe = pos_src.flatten(2).permute(0, 2, 1) |
| 341 | queries = tokens |
| 342 | query_pe = tokens |
| 343 | |
| 344 | for idx in range(2): |
| 345 | prefix = f"transformer.layers.{idx}" |
| 346 | if idx == 0: |
| 347 | queries = attention_forward( |
| 348 | queries, queries, queries, prefix + ".self_attn", mask_weights, 8 |
| 349 | ) |
| 350 | else: |
| 351 | q = queries + query_pe |
| 352 | attn_out = attention_forward(q, q, queries, prefix + ".self_attn", mask_weights, 8) |
| 353 | queries = queries + attn_out |
| 354 | queries = F.layer_norm( |
| 355 | queries, |
| 356 | [dim], |
| 357 | mask_weights[prefix + ".norm1.weight"].float(), |
| 358 | mask_weights[prefix + ".norm1.bias"].float(), |
| 359 | ) |
| 360 | |
| 361 | q = queries + query_pe |
| 362 | k = keys + key_pe |
| 363 | attn_out = attention_forward( |
no test coverage detected