paligemma2-10b-pt-448-adapter_detection_coco

Fine-tuning PaliGemma-2 for closed-set object detection on the COCO2017 dataset.

See Technical Report for details.

Model Loading

from transformers import (
    PaliGemmaProcessor,
    PaliGemmaForConditionalGeneration,
)
from peft import PeftModel
from transformers.image_utils import load_image
import torch

url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg"
image = load_image(url)

model_id = "google/paligemma2-10b-pt-448"
adapter_id = "segis95/paligemma2-10b-pt-448-adapter_detection_coco"

base_model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto").eval()
model = PeftModel.from_pretrained(base_model, adapter_id).eval()
processor = PaliGemmaProcessor.from_pretrained(model_id, use_fast=True)

Basic Usage

prompt = "<image>detect all classes\n"
model_inputs = processor(text=prompt, images=image, return_tensors="pt").to(torch.bfloat16).to(model.device)
input_len = model_inputs["input_ids"].shape[-1]

with torch.inference_mode():
    generation = model.generate(**model_inputs, max_new_tokens=256, do_sample=False)
    generation = generation[0][input_len:]
    decoded = processor.decode(generation, skip_special_tokens=True)
    print(decoded)

Restricted Decoding

Exactly 256 new tokens are generated, consisting of 51 groups (representing bounding box predictions) following the pattern <locAAAA><locBBBB><locCCCC><locDDDD><unusedXX>, followed by a final <eos> token.

Every 5th generated token (of type <unusedXX>, where XX ranges from 0 to 80) represents either one of the 80 COCO classes or a 'noise' token. The noise token explicitly represents 'no class present' and is intentionally weighted to dominate over actual class tokens.

Each noise token (token ID 87, corresponding to <unused80>) will be subsequently replaced with the second most probable class token (token IDs ranging from 7 (<unused0>) to 86 (<unused79>)). The probability scores of these class tokens are used as confidence scores for the predicted 51 bounding boxes.

Restricted decoding without noise token replacement

from transformers import LogitsProcessor, LogitsProcessorList
import torch.nn.functional as F

class PositionBasedTokenFilter(LogitsProcessor):
    def __init__(self, prefix_len):
        self.prefix_len = prefix_len
        self.first_class_token = 7
        self.noise_token = 87
        self.first_loc_token = 256000
        self.last_loc_token = 257023

    def __call__(
        self, input_ids: torch.LongTensor, scores: torch.FloatTensor
    ) -> torch.FloatTensor:
        """
        Generating exactly 256 new tokens.
        There are exactly 51 groups of type <locAAAA><locBBBB><locCCCC><locDDDD><unusedXX> followed by <eos> at the end.
        """
        
        current_position = input_ids.shape[-1] - self.prefix_len
        mask = torch.full_like(scores, float("-inf"))
        if current_position == 255:
            mask[:, 1] = 0.0 # <eos>
        elif (current_position + 1) % 5 != 0:
            mask[:, self.first_loc_token: self.last_loc_token + 1] = 0.0
        else:
            mask[:, self.first_class_token: self.noise_token + 1] = 0.0

        scores = scores + mask

        return scores

with torch.inference_mode():
    generation = model.generate(
        **model_inputs,
        max_new_tokens=256,
        do_sample=False,
        return_dict_in_generate=True,
        output_scores=True,
        logits_processor=LogitsProcessorList(
            [PositionBasedTokenFilter(input_len)]
        ),
    )

    generation = generation['sequences'][0, input_len:]
    decoded = processor.decode(generation, skip_special_tokens=True)
    print(decoded)

Restricted decoding with noise token replacement

def generate_and_replace_noise_with_class_tokens(model,
                                               processor,
                                               model_inputs,
                                               input_len,
                                               new_tokens=256,
                                               noise_token=87
                                              ):

    with torch.inference_mode():
        generation = model.generate(
            **model_inputs,
            max_new_tokens=new_tokens,
            do_sample=False,
            return_dict_in_generate=True,
            output_scores=True,
            logits_processor=LogitsProcessorList(
                [PositionBasedTokenFilter(input_len)]
            ),
        )
    
    sequence = (
        generation.sequences[0, input_len:].detach().cpu().numpy().tolist()
    )

    scores = torch.stack(
        [
            generation.scores[pos_id][0].detach().cpu()
            for pos_id in range(new_tokens)
        ]
    )
    probabilities = F.softmax(scores, dim=-1)

    assert len(sequence) == len(probabilities) == new_tokens

    result = []
    for token, probas in zip(sequence, probabilities):
        p_top, tokens_top = torch.topk(probas, 2)
    
        assert torch.isclose(probas[token], probas[tokens_top[0]])
    
        if noise_token == tokens_top[0]:
            result.append(
                (
                    processor.tokenizer.decode(tokens_top[1].item()),
                    p_top[1].item(),
                )
            )
        else:
            result.append(
                (
                    processor.tokenizer.decode(tokens_top[0].item()),
                    p_top[0].item(),
                )
            )
    return result

tokens_with_scores = generate_and_replace_noise_with_class_tokens(model, processor, model_inputs, input_len)
print(tokens_with_scores)