Refactoring to distangle modules

src/config.py

@@ -65,7 +65,7 @@ class TinyCLIPConfig(PretrainedConfig):
         max_len: int = 128,
         cls_type: bool = True,
         freeze_vision_base: bool = False,
-        freeze_text_base: bool = False,
+        freeze_text_base: bool = True,
         loss_type: str = "cyclip",
         **kwargs,
     ):
@@ -85,18 +85,6 @@ class TinyCLIPConfig(PretrainedConfig):
         super().__init__(**kwargs)
 
 
-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 = 64
@@ -112,5 +100,5 @@ class TrainerConfig(pydantic.BaseModel):
 
     run_openai_clip: bool = False
 
-    _model_config: ModelConfig = ModelConfig()
+    _model_config: TinyCLIPConfig = TinyCLIPConfig()
     _data_config: DataConfig = DataConfig()

src/models.py

@@ -1,14 +1,14 @@
 from PIL import Image
-import timm
-from timm import data
+import transformers
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import transformers
+
 from transformers import PreTrainedModel
 
 from src.config import TinyCLIPConfig, TinyCLIPTextConfig, TinyCLIPVisionConfig
 from src import loss
+from src import vision_model
 
 
 class Projection(nn.Module):
@@ -70,9 +70,10 @@ class TinyCLIPVisionEncoder(PreTrainedModel):
 
     def __init__(self, config: TinyCLIPVisionConfig):
         super().__init__(config)
-
+        base, num_features = vision_model.get_vision_base(config)
+        self.base = base
         self.projection = projection_layers(
-            self.base.num_features, config.embed_dims, config.projection_layers
+            num_features, config.embed_dims, config.projection_layers
         )
 
     def forward(self, images: list[Image.Image]):

src/tokenizer.py

@@ -3,11 +3,13 @@ from typing import Union
 import torch
 from transformers import AutoTokenizer
 
+from src.config import TinyCLIPTextConfig
+
 
 class Tokenizer:
-    def __init__(self, model_name: str, max_len: int) -> None:
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.max_len = max_len
+    def __init__(self, text_config: TinyCLIPTextConfig) -> None:
+        self.tokenizer = AutoTokenizer.from_pretrained(text_config.text_model)
+        self.max_len = text_config.max_len
 
     def __call__(self, x: Union[str, list[str]]) -> dict[str, torch.LongTensor]:
         return self.tokenizer(

src/trainer.py

@@ -1,7 +1,25 @@
 from src import data
 from src import config
 from src import vision_model
+from src import tokenizer as tk
+from src.lightning_module import LightningModule
+from src import loss
+from src import models
 
 
 def train(config: config.TrainerConfig):
-    train_dl, valid_dl = data.get_dataset()
+    transform = vision_model.get_vision_transform(config._model_config.vision_config)
+    tokenizer = tk.Tokenizer(config._model_config.text_config)
+    train_dl, valid_dl = data.get_dataset(
+        transform=transform, tokenizer=tokenizer, hyper_parameters=config  # type: ignore
+    )
+    vision_encoder = models.TinyCLIPVisionEncoder(config=config._model_config.vision_config)
+    text_encoder = models.TinyCLIPTextEncoder(config=config._model_config.text_config)
+
+    lightning_module = LightningModule(
+        vision_encoder=vision_encoder,
+        text_encoder=text_encoder,
+        loss_fn=loss.get_loss(config._model_config.loss_type),
+        hyper_parameters=config,
+        len_train_dl=len(train_dl),
+    )

src/vision_model.py

@@ -1,11 +1,20 @@
 import timm
 from timm import data
+import torch.nn as nn
+from torchvision import transforms
 
-from src import config
+from src.config import TinyCLIPVisionConfig
 
 
-def get_vision_base_and_transform(config: config.TrainerConfig):
-    base = timm.create_model(config._model_config.vision_model, num_classes=0)
-    timm_config = data.resolve_data_config({}, model=base)
+def get_vision_base(
+    config: TinyCLIPVisionConfig,
+) -> tuple[nn.Module, int]:
+    base = timm.create_model(config.vision_model, num_classes=0, pretrained=True)
+    num_features = base.num_features
+    return base, num_features
+
+
+def get_vision_transform(config: TinyCLIPVisionConfig) -> transforms.Compose:
+    timm_config = data.resolve_data_config({}, model=config.vision_model)
     transform = data.transforms_factory.create_transform(**timm_config)
-    return base, transform
+    return transform  # type: ignore