Remove lightning dependency from submission

evaluate/data.py

@@ -1,10 +1,9 @@
 from pathlib import Path
 
 import pandas as pd
-import pytorch_lightning as pl
 import torch
 from PIL import Image
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import Dataset
 from torchvision.transforms import v2
 from transformers import AutoImageProcessor, AutoModel
 
@@ -40,30 +39,3 @@ class ImageDataset(Dataset):
         img = Image.open(image_path).convert("RGB")
         img = v2.ToTensor()(img)
         return {"features": img, "observation_id": row.observation_id}
-
-
-class InferenceDataModel(pl.LightningDataModule):
-    def __init__(
-        self,
-        metadata_path,
-        images_root_path,
-        batch_size=32,
-    ):
-        super().__init__()
-        self.metadata_path = metadata_path
-        self.images_root_path = images_root_path
-        self.batch_size = batch_size
-
-    def setup(self, stage=None):
-        self.dataloader = DataLoader(
-            ImageDataset(self.metadata_path, self.images_root_path),
-            batch_size=self.batch_size,
-            shuffle=False,
-            num_workers=4,
-        )
-
-    def predict_dataloader(self):
-        transform = v2.Compose([TransformDino("facebook/dinov2-base")])
-        for batch in self.dataloader:
-            batch = transform(batch)
-            yield batch

evaluate/data_lightning.py

@@ -0,0 +1,31 @@
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from torchvision.transforms import v2
+from .data import ImageDataset, TransformDino
+
+
+class InferenceDataModel(pl.LightningDataModule):
+    def __init__(
+        self,
+        metadata_path,
+        images_root_path,
+        batch_size=32,
+    ):
+        super().__init__()
+        self.metadata_path = metadata_path
+        self.images_root_path = images_root_path
+        self.batch_size = batch_size
+
+    def setup(self, stage=None):
+        self.dataloader = DataLoader(
+            ImageDataset(self.metadata_path, self.images_root_path),
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=4,
+        )
+
+    def predict_dataloader(self):
+        transform = v2.Compose([TransformDino("facebook/dinov2-base")])
+        for batch in self.dataloader:
+            batch = transform(batch)
+            yield batch

evaluate/submission.py

@@ -1,11 +1,21 @@
-import zipfile
-
 import pandas as pd
 import torch
-from pytorch_lightning import Trainer
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision.transforms import v2
+
+from .data import ImageDataset, TransformDino
+
 
-from .data import InferenceDataModel
-from .model import LinearClassifier
+class LinearClassifier(nn.Module):
+    def __init__(self, num_features, num_classes):
+        super().__init__()
+        self.num_features = num_features
+        self.num_classes = num_classes
+        self.model = nn.Linear(num_features, num_classes)
+
+    def forward(self, x):
+        return torch.log_softmax(self.model(x), dim=1)
 
 
 def make_submission(
@@ -14,17 +24,22 @@ def make_submission(
     output_csv_path="./submission.csv",
     images_root_path="/tmp/data/private_testset",
 ):
-    model = LinearClassifier.load_from_checkpoint(model_path)
-    dm = InferenceDataModel(
-        metadata_path=test_metadata, images_root_path=images_root_path
-    )
-    trainer = Trainer(
-        accelerator="gpu" if torch.cuda.is_available() else "cpu",
+    checkpoint = torch.load(model_path)
+    hparams = checkpoint["hyper_parameters"]
+    model = LinearClassifier(hparams["num_features"], hparams["num_classes"])
+    model.load_state_dict(checkpoint["state_dict"])
+
+    transform = v2.Compose([TransformDino("facebook/dinov2-base")])
+    dataloader = DataLoader(
+        ImageDataset(test_metadata, images_root_path), batch_size=32, num_workers=4
     )
-    predictions = trainer.predict(model, datamodule=dm)
     rows = []
-    for batch in predictions:
-        for observation_id, class_id in zip(batch["observation_id"], batch["class_id"]):
+    for batch in dataloader:
+        batch = transform(batch)
+        observation_ids = batch["observation_id"]
+        logits = model(batch["features"])
+        class_ids = torch.argmax(logits, dim=1)
+        for observation_id, class_id in zip(observation_ids, class_ids):
             row = {"observation_id": int(observation_id), "class_id": int(class_id)}
             rows.append(row)
     submission_df = pd.DataFrame(rows)

evaluate/test_evaluate.py

@@ -5,8 +5,9 @@ import pytest
 import torch
 from pytorch_lightning import Trainer
 
-from .data import ImageDataset, InferenceDataModel
-from .model import LinearClassifier
+from .data import ImageDataset
+from .data_lightning import InferenceDataModel
+from .model_lightning import LinearClassifier as LightningLinearClassifier
 from .submission import make_submission
 
 
@@ -43,7 +44,7 @@ def metadata(tmp_path, images_root):
 @pytest.fixture
 def model_checkpoint(tmp_path, metadata, images_root):
     model_checkpoint = tmp_path / "model.ckpt"
-    model = LinearClassifier(768, 10)
+    model = LightningLinearClassifier(768, 10)
     trainer = Trainer(max_epochs=1, fast_dev_run=True)
     dm = TestingInferenceDataModel(metadata, images_root)
     trainer.fit(model, dm)
@@ -69,7 +70,7 @@ def test_inference_datamodel(images_root, metadata):
 
 
 def test_model_checkpoint(model_checkpoint):
-    model = LinearClassifier.load_from_checkpoint(model_checkpoint)
+    model = LightningLinearClassifier.load_from_checkpoint(model_checkpoint)
     assert model