timm

/

vit_large_patch14_clip_224.openai_ft_in12k_in1k

like 36

Follow

PyTorch Image Models 140

You can create your own like this:

            feature_extractor = CustomImageProcessor(224)
            model = ViTForImageClassification.from_pretrained("timm/vit_large_patch14_clip_224.openai_ft_in12k_in1k",
                                                              num_labels=num_classes)

where:

class CustomImageProcessor(ViTImageProcessor):
    def __init__(self, size):
        super().__init__(size=size)

    def __call__(self, images, return_tensors="pt"):
        return super().__call__(images=images, return_tensors=return_tensors)

and

class ImageDataCollator:
    def __init__(self, feature_extractor):
        self.feature_extractor = feature_extractor

    def __call__(self, batch):
        images, labels = zip(*batch)
        inputs = self.feature_extractor(images=images, return_tensors="pt")
        inputs["labels"] = torch.tensor(labels, dtype=torch.long)
        return inputs

then

            data_collator = ImageDataCollator(feature_extractor)

            trainer = Trainer(
                model=model,
                args=training_args,
                train_dataset=train_dataset,
                eval_dataset=val_dataset,
                data_collator=data_collator,
            )

FYI this is not a transformers ViT model, it's a timm model