File size: 1,189 Bytes
5e1c8df 35352c6 5e1c8df 35352c6 5e1c8df 35352c6 5e1c8df 35352c6 5e1c8df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import pathlib
import pydantic
MAX_DOWNLOAD_TIME = 0.2
IMAGE_DOWNLOAD_PATH = pathlib.Path("/tmp/images")
class DataConfig(pydantic.BaseModel):
buffer_size: int = 1000
data_len: int = 100
train_len: int = 90
small_dataset: str = "laion/220k-gpt4vision-captions-from-livis"
large_dataset: str = "laion/laion400m"
dataset: str = small_dataset
class ModelConfig(pydantic.BaseModel):
text_model: str = "microsoft/xtremedistil-l6-h256-uncased" # 51 mb
vision_model: str = "edgenext_small" # 20 mb
projection_layers: int = 3
embed_dim: int = 256
transformer_embed_dim: int = 768
max_len: int = 128 # 77
cls_type: bool = True
freeze_vision_base: bool = False
freeze_text_base: bool = False
class TrainerConfig(pydantic.BaseModel):
epochs: int = 20
batch_size: int = 256
learning_rate: float = 5e-4
accumulate_grad_batches: int = 1
temperature: float = 1.0
vision_freeze_layers: int = 2
lambda_1: float = 1.0
lambda_2: float = 1.0
val_check_interval: int = 1000
run_openai_clip: bool = False
_model_config: ModelConfig = ModelConfig()
_data_config: DataConfig = DataConfig()
|