ColossalAI/tests/kit/model_zoo/transformers/sam.py

import torch
import transformers

from ..registry import ModelAttribute, model_zoo

# ===============================
# Register single-image SAM
# ===============================


# define data gen function
def data_gen():
    # Generated from following code snippet
    #
    # from PIL import Image
    # import requests
    # from transformers import SamModel, SamProcessor
    #
    # model = SamModel.from_pretrained("facebook/sam-vit-base")
    # processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
    #
    # img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
    # raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
    # input_points = [[[450, 600]]] # 2D localization of a window
    # inputs = processor(raw_image, input_points=input_points, return_tensors="pt")

    pixel_values = torch.rand(1, 3, 1024, 1024, dtype=torch.float32)
    original_sizes = torch.tensor([[1764, 2646]], dtype=torch.int64)
    reshaped_input_sizes = torch.tensor([[683, 1024]], dtype=torch.int64)
    input_points = torch.tensor([[[[174.1497, 232.3129]]]], dtype=torch.float64)
    return dict(
        pixel_values=pixel_values,
        original_sizes=original_sizes,
        reshaped_input_sizes=reshaped_input_sizes,
        input_points=input_points,
    )


# define output transform function
output_transform_fn = lambda x: x

# define loss funciton
loss_fn = lambda x: x["iou_scores"].mean()

config = transformers.SamConfig()
config.vision_config.num_hidden_layers = 2

# register the BERT variants
model_zoo.register(
    name="transformers_sam",
    model_fn=lambda: transformers.SamModel(config),
    data_gen_fn=data_gen,
    output_transform_fn=output_transform_fn,
    loss_fn=loss_fn,
    model_attribute=ModelAttribute(has_control_flow=True),
)