succesful local run

.vscode/settings.json

@@ -2,7 +2,7 @@
     "files.insertFinalNewline": true,
     "jupyter.debugJustMyCode": false,
     "editor.formatOnSave": true,
-    "editor.formatOnPaste": true,
+    // "editor.formatOnPaste": true,
     "files.autoSave": "onFocusChange",
     "editor.defaultFormatter": "ms-python.black-formatter",
     "black-formatter.path": ["/opt/homebrew/bin/black"],
@@ -12,12 +12,12 @@
     "isort.check": true,
     "python.analysis.typeCheckingMode": "basic",
     "python.defaultInterpreterPath": "/opt/homebrew/bin/python3",
-    "[python]": {
-        "editor.defaultFormatter": "ms-python.black-formatter",
-        "editor.formatOnSave": true,
-        "editor.codeActionsOnSave": {
-            "source.organizeImports": "explicit"
-        },
-    },
-    "isort.args":["--profile", "black"],
+    // "[python]": {
+    //     "editor.defaultFormatter": "ms-python.black-formatter",
+    //     "editor.formatOnSave": true,
+    //     "editor.codeActionsOnSave": {
+    //         "source.organizeImports": "explicit"
+    //     },
+    // },
+    // "isort.args":["--profile", "black"],
 }

src/config.py

@@ -6,6 +6,14 @@ from transformers import PretrainedConfig
 MAX_DOWNLOAD_TIME = 0.2
 
 IMAGE_DOWNLOAD_PATH = pathlib.Path("./data/images")
+WANDB_LOG_PATH = pathlib.Path("/tmp/wandb_logs")
+
+IMAGE_DOWNLOAD_PATH.mkdir(parents=True, exist_ok=True)
+WANDB_LOG_PATH.mkdir(parents=True, exist_ok=True)
+
+MODEL_NAME = "tiny_clip"
+
+WANDB_ENTITY = "sachinruk"
 
 
 class DataConfig(pydantic.BaseModel):
@@ -97,6 +105,8 @@ class TrainerConfig(pydantic.BaseModel):
     lambda_2: float = 1.0
 
     val_check_interval: int = 1000
+    log_every_n_steps: int = 100
+    debug: bool = False
 
     run_openai_clip: bool = False
 

src/data.py

@@ -37,7 +37,7 @@ class CollateFn:
         tokenized_text = self.tokenizer([item["caption"] for item in batch])
 
         return {
-            "image": stacked_images,
+            "images": stacked_images,
             **tokenized_text,
         }
 

src/lightning_module.py

@@ -24,10 +24,11 @@ class LightningModule(pl.LightningModule):
         self.hyper_parameters = hyper_parameters
         self.len_train_dl = len_train_dl
 
-    def common_step(self, batch: tuple[torch.Tensor, list[str]], step_kind: str) -> torch.Tensor:
-        text, images = batch
-        image_features = self.vision_encoder(images)
-        text_features = self.text_encoder(text)
+    def common_step(self, batch: dict[str, torch.Tensor], step_kind: str) -> torch.Tensor:
+        image_features = self.vision_encoder(batch["images"])
+        text_features = self.text_encoder(
+            {key: value for key, value in batch.items() if key != "images"}
+        )
         similarity_matrix = loss_utils.get_similarity_matrix(image_features, text_features)
 
         loss = self.loss_fn(similarity_matrix, image_features, text_features)
@@ -52,10 +53,6 @@ class LightningModule(pl.LightningModule):
                 "params": self.vision_encoder.projection.parameters(),
                 "lr": self.hyper_parameters.learning_rate,
             },
-            {
-                "params": self.vision_encoder.base.parameters(),
-                "lr": self.hyper_parameters.learning_rate / 2,
-            },
         ]
         caption_params = [
             {

src/models.py

@@ -77,10 +77,8 @@ class TinyCLIPVisionEncoder(PreTrainedModel):
             num_features, config.embed_dims, config.projection_layers
         )
 
-    def forward(self, images: list[Image.Image]):
-        x: torch.Tensor = torch.stack([self.transform(image) for image in images])  # type: ignore
-
-        projected_vec = self.projection(self.base(x))
+    def forward(self, images: torch.Tensor):
+        projected_vec = self.projection(self.base(images))
         return F.normalize(projected_vec, dim=-1)
 
 

src/trainer.py

@@ -1,25 +1,34 @@
-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 data
 from src import loss
 from src import models
+from src import tokenizer as tk
+from src import vision_model
+from src import utils
+from src.lightning_module import LightningModule
 
 
-def train(config: config.TrainerConfig):
-    transform = vision_model.get_vision_transform(config._model_config.vision_config)
-    tokenizer = tk.Tokenizer(config._model_config.text_config)
+def train(trainer_config: config.TrainerConfig):
+    transform = vision_model.get_vision_transform(trainer_config._model_config.vision_config)
+    tokenizer = tk.Tokenizer(trainer_config._model_config.text_config)
     train_dl, valid_dl = data.get_dataset(
-        transform=transform, tokenizer=tokenizer, hyper_parameters=config  # type: ignore
+        transform=transform, tokenizer=tokenizer, hyper_parameters=trainer_config  # type: ignore
     )
-    vision_encoder = models.TinyCLIPVisionEncoder(config=config._model_config.vision_config)
-    text_encoder = models.TinyCLIPTextEncoder(config=config._model_config.text_config)
+    vision_encoder = models.TinyCLIPVisionEncoder(config=trainer_config._model_config.vision_config)
+    text_encoder = models.TinyCLIPTextEncoder(config=trainer_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,
+        loss_fn=loss.get_loss(trainer_config._model_config.loss_type),
+        hyper_parameters=trainer_config,
         len_train_dl=len(train_dl),
     )
+
+    trainer = utils.get_trainer(trainer_config)
+    trainer.fit(lightning_module, train_dl, valid_dl)
+
+
+if __name__ == "__main__":
+    trainer_config = config.TrainerConfig(debug=True)
+    train(trainer_config)

src/utils.py

@@ -0,0 +1,31 @@
+import datetime
+
+import pytorch_lightning as pl
+from pytorch_lightning import loggers
+
+from src import config
+
+
+def _get_wandb_logger(trainer_config: config.TrainerConfig):
+    name = f"{config.MODEL_NAME}-{datetime.datetime.now()}"
+    if trainer_config.debug:
+        name = "debug-" + name
+    return loggers.WandbLogger(
+        entity=config.WANDB_ENTITY,
+        save_dir=config.WANDB_LOG_PATH,
+        project=config.MODEL_NAME,
+        name=name,
+        config=trainer_config._model_config.to_dict(),
+    )
+
+
+def get_trainer(trainer_config: config.TrainerConfig):
+    return pl.Trainer(
+        max_epochs=trainer_config.epochs if not trainer_config.debug else 1,
+        logger=_get_wandb_logger(trainer_config),
+        log_every_n_steps=trainer_config.log_every_n_steps,
+        gradient_clip_val=1.0,
+        limit_train_batches=5 if trainer_config.debug else 1.0,
+        limit_val_batches=5 if trainer_config.debug else 1.0,
+        accelerator="auto",
+    )