|
""" ViTamin |
|
|
|
Paper: Designing Scalable Vison Models in the Vision-Language Era |
|
|
|
@misc{chen2023designing, |
|
title={Designing Scalable Vison Models in the Vision-Language Era}, |
|
author={Jieneng Chen and Qihang Yu and Xiaohui Shen and Alan Yuille and Liang-Cheih Chen}, |
|
year={2023}, |
|
archivePrefix={arXiv}, |
|
primaryClass={cs.CV} |
|
} |
|
|
|
Based on Apache 2.0 licensed code at https://github.com/Beckschen/ViTamin |
|
|
|
by Jieneng Chen 2024 |
|
""" |
|
|
|
import copy |
|
import os |
|
from collections import OrderedDict |
|
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union |
|
|
|
|
|
if TYPE_CHECKING: |
|
from transformers.processing_utils import ProcessorMixin |
|
from transformers.utils import TensorType |
|
|
|
from transformers.configuration_utils import PretrainedConfig |
|
from transformers.utils import logging |
|
|
|
logger = logging.get_logger(__name__) |
|
|
|
class ViTaminTextConfig(PretrainedConfig): |
|
model_type = "vitamin_text_model" |
|
|
|
def __init__( |
|
self, |
|
context_length = 77, |
|
vocab_size = 49408, |
|
width = 1024, |
|
heads = 16, |
|
layers = 24, |
|
**kwargs, |
|
): |
|
super().__init__(**kwargs) |
|
|
|
self.vocab_size = vocab_size |
|
self.context_length = context_length |
|
self.width = width |
|
self.heads = heads |
|
self.layers = layers |
|
|
|
@classmethod |
|
def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig": |
|
config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs) |
|
|
|
if 'text_config' in config_dict: |
|
config_dict = config_dict['text_config'] |
|
|
|
if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type: |
|
logger.warning( |
|
f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " |
|
f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." |
|
) |
|
|
|
return cls.from_dict(config_dict, **kwargs) |
|
|
|
|
|
class ViTaminVisionConfig(PretrainedConfig): |
|
|
|
model_type = "vitamin_vision_model" |
|
|
|
def __init__( |
|
self, |
|
timm_model_name = "vitamin_large", |
|
timm_model_pretrained = False, |
|
timm_pool = "", |
|
timm_proj = "linear", |
|
timm_drop = 0.0, |
|
timm_drop_path = 0.1, |
|
image_size = 256, |
|
timm_proj_bias = False, |
|
patch_dropout = 0.0, |
|
drop_path = None, |
|
**kwargs, |
|
): |
|
super().__init__(**kwargs) |
|
|
|
self.timm_model_name = timm_model_name |
|
self.timm_model_pretrained = timm_model_pretrained |
|
self.timm_pool = timm_pool |
|
self.timm_proj = timm_proj |
|
self.timm_drop = timm_drop |
|
self.timm_drop_path = timm_drop_path |
|
self.timm_proj_bias = timm_proj_bias |
|
self.patch_dropout = patch_dropout |
|
self.image_size = image_size |
|
|
|
|
|
@classmethod |
|
def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig": |
|
config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs) |
|
|
|
if 'vision_config' in config_dict: |
|
config_dict = config_dict['vision_config'] |
|
|
|
if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type: |
|
logger.warning( |
|
f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " |
|
f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." |
|
) |
|
|
|
return cls.from_dict(config_dict, **kwargs) |
|
|
|
|
|
|
|
class ViTaminConfig(PretrainedConfig): |
|
model_type = "vitamin" |
|
is_composition = True |
|
|
|
def __init__( |
|
self, text_config=None, vision_config=None, embed_dim=512, **kwargs |
|
): |
|
super().__init__(**kwargs) |
|
if text_config is None: |
|
text_config = {} |
|
logger.info("`text_config` is `None`. Initializing the `CLIPTextConfig` with default values.") |
|
|
|
if vision_config is None: |
|
vision_config = {} |
|
logger.info("`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values.") |
|
|
|
self.embed_dim = embed_dim |
|
self.text_config = ViTaminTextConfig(**text_config) |
|
self.vision_config = ViTaminVisionConfig(**vision_config) |
|
|
|
@classmethod |
|
def from_text_vision_configs(cls, text_config: ViTaminTextConfig, vision_config: ViTaminVisionConfig, **kwargs): |
|
r""" |
|
Instantiate a [`CLIPConfig`] (or a derived class) from clip text model configuration and clip vision model |
|
configuration. |
|
Returns: |
|
[`CLIPConfig`]: An instance of a configuration object |
|
""" |
|
|
|
return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs) |
|
|
|
def to_dict(self): |
|
""" |
|
Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. |
|
Returns: |
|
`Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance, |
|
""" |
|
output = copy.deepcopy(self.__dict__) |
|
output["text_config"] = self.text_config.to_dict() |
|
output["vision_config"] = self.vision_config.to_dict() |
|
output["model_type"] = self.__class__.model_type |
|
return output |
|
|
