Function dump_case

tests/dump_phase6_reference.py:256–441

(case: PromptCase,
              prompt_weights: Dict[str, torch.Tensor],
              mask_weights: Dict[str, torch.Tensor],
              no_mem_embed: torch.Tensor,
              neck_trk_0: torch.Tensor,
              neck_trk_1: torch.Tensor,
              neck_trk_2: torch.Tensor,
              outdir: str)

def dump_case(case: PromptCase,
              prompt_weights: Dict[str, torch.Tensor],
              mask_weights: Dict[str, torch.Tensor],
              no_mem_embed: torch.Tensor,
              neck_trk_0: torch.Tensor,
              neck_trk_1: torch.Tensor,
              neck_trk_2: torch.Tensor,
              outdir: str) -> None:
    os.makedirs(outdir, exist_ok=True)

    sparse_embeddings, dense_embeddings, image_pe = build_prompt_encoder_outputs(prompt_weights, case)

    image_embeddings = neck_trk_2 + no_mem_embed.view(1, -1, 1, 1)
    feat_s0 = F.conv2d(
        neck_trk_0,
        mask_weights["conv_s0.weight"].float(),
        mask_weights["conv_s0.bias"].float(),
    )
    feat_s1 = F.conv2d(
        neck_trk_1,
        mask_weights["conv_s1.weight"].float(),
        mask_weights["conv_s1.bias"].float(),
    )

    save_ggml_bnd(os.path.join(outdir, "sam_pe_sparse"), sparse_embeddings)
    save_ggml_nchw(os.path.join(outdir, "sam_pe_dense"), dense_embeddings)
    save_ggml_nchw(os.path.join(outdir, "sam_pe_image_pe"), image_pe)
    save_ggml_nchw(os.path.join(outdir, "sam_dec_image_feats"), image_embeddings)

    output_tokens = torch.cat(
        [
            mask_weights["obj_score_token.weight"].float(),
            mask_weights["iou_token.weight"].float(),
            mask_weights["mask_tokens.weight"].float(),
        ],
        dim=0,
    ).unsqueeze(0).expand(sparse_embeddings.size(0), -1, -1)
    tokens = torch.cat((output_tokens, sparse_embeddings), dim=1)
    save_ggml_bnd(os.path.join(outdir, "sam_dec_tokens_initial"), tokens)

    src = image_embeddings + dense_embeddings
    pos_src = image_pe.repeat(tokens.shape[0], 1, 1, 1)
    bsz, dim, h, w = src.shape
    keys = src.flatten(2).permute(0, 2, 1)
    key_pe = pos_src.flatten(2).permute(0, 2, 1)
    queries = tokens
    query_pe = tokens

    for idx in range(2):
        prefix = f"transformer.layers.{idx}"
        if idx == 0:
            queries = attention_forward(
                queries, queries, queries, prefix + ".self_attn", mask_weights, 8
            )
        else:
            q = queries + query_pe
            attn_out = attention_forward(q, q, queries, prefix + ".self_attn", mask_weights, 8)
            queries = queries + attn_out