ViTamin-XL-384px / configuration_vitamin.py
bbexx's picture
add dependency
16d57ae
""" 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