Introduce CoresetMethodsSupport

modyn/models/README.md

@@ -5,4 +5,9 @@ The user can define models here. The model definition should take as a parameter
 # Wild Time models
 The code for the models used for WildTime is taken from the official [repository](https://github.com/huaxiuyao/Wild-Time). 
 The original version is linked in each class.
-You can find [here](https://raw.githubusercontent.com/huaxiuyao/Wild-Time/main/LICENSE) a copy of the MIT license
+You can find [here](https://raw.githubusercontent.com/huaxiuyao/Wild-Time/main/LICENSE) a copy of the MIT license
+
+# Embedding Recorder
+Many coreset methods are adapted from the [DeepCore](https://github.com/PatrickZH/DeepCore/) library. To use them, the models must keep track of the embeddings (activations of the penultimate layer). This is
+done using the `EmbeddingRecorder` class, which is adapted from the aforementioned project.
+You can find a copy of their MIT license [here](https://raw.githubusercontent.com/PatrickZH/DeepCore/main/LICENSE.md)

modyn/models/articlenet/articlenet.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 import torch
+from modyn.models.modyn_model import ModynModel
 from torch import nn
 from transformers import DistilBertModel
 
@@ -43,12 +44,16 @@ def __call__(self, data: torch.Tensor) -> torch.Tensor:
         return pooled_output
 
 
-class ArticleNetwork(nn.Module):
+class ArticleNetwork(ModynModel):
     def __init__(self, num_classes: int) -> None:
         super().__init__()
         self.featurizer = DistilBertFeaturizer.from_pretrained("distilbert-base-uncased")
         self.classifier = nn.Linear(self.featurizer.d_out, num_classes)
 
     def forward(self, data: torch.Tensor) -> torch.Tensor:
         embedding = self.featurizer(data)
+        embedding = self.embedding_recorder(embedding)
         return self.classifier(embedding)
+
+    def get_last_layer(self) -> nn.Module:
+        return self.classifier

modyn/models/dlrm/dlrm.py

@@ -1,11 +1,12 @@
-from typing import Any
+from typing import Any, Optional
 
 import numpy as np
 import torch
 from modyn.models.dlrm.nn.factories import create_interaction
 from modyn.models.dlrm.nn.parts import DlrmBottom, DlrmTop
 from modyn.models.dlrm.utils.install_lib import install_cuda_extensions_if_not_present
 from modyn.models.dlrm.utils.utils import get_device_mapping
+from modyn.models.modyn_model import EmbeddingRecorder, ModynModel
 from modyn.utils.utils import package_available_and_can_be_imported
 from torch import nn
 
@@ -16,7 +17,7 @@ def __init__(self, model_configuration: dict[str, Any], device: str, amp: bool)
         self.model.to(device)
 
 
-class DlrmModel(nn.Module):
+class DlrmModel(ModynModel):
     # pylint: disable=too-many-instance-attributes
     def __init__(self, model_configuration: dict[str, Any], device: str, amp: bool) -> None:
         super().__init__()
@@ -124,3 +125,15 @@ def forward(self, data: torch.Tensor) -> torch.Tensor:
             numerical_input, self.reorder_categorical_input(categorical_input)
         )
         return self.top_model(from_bottom, bottom_mlp_output).squeeze()
+
+    # delegate the embedding handling to the top model
+    @property
+    def embedding(self) -> Optional[torch.Tensor]:
+        return self.top_model.embedding
+
+    @property
+    def embedding_recorder(self) -> EmbeddingRecorder:
+        return self.top_model.embedding_recorder
+
+    def get_last_layer(self) -> nn.Module:
+        return self.top_model.get_last_layer()

modyn/models/dlrm/nn/parts.py

@@ -22,6 +22,7 @@
 from modyn.models.dlrm.nn.embeddings import Embeddings
 from modyn.models.dlrm.nn.factories import create_embeddings, create_mlp
 from modyn.models.dlrm.nn.interactions import Interaction
+from modyn.models.modyn_model import ModynModel
 from torch import nn
 
 
@@ -104,7 +105,7 @@ def forward(self, numerical_input, categorical_inputs) -> Tuple[torch.Tensor, Op
         return torch.cat(bottom_output, dim=1), bottom_mlp_output
 
 
-class DlrmTop(nn.Module):
+class DlrmTop(ModynModel):
     def __init__(self, top_mlp_sizes: Sequence[int], interaction: Interaction, device: str, use_cpp_mlp: bool = False):
         super().__init__()
 
@@ -127,4 +128,9 @@ def forward(self, bottom_output, bottom_mlp_output):
             bottom_mlp_output (Tensor): with shape [batch_size, embedding_dim]
         """
         interaction_output = self.interaction.interact(bottom_output, bottom_mlp_output)
-        return self.out(self.mlp(interaction_output))
+        mlp_output = self.mlp(interaction_output)
+        mlp_output = self.embedding_recorder(mlp_output)
+        return self.out(mlp_output)
+
+    def get_last_layer(self) -> nn.Module:
+        return self.out

modyn/models/fmownet/fmownet.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.nn.functional as F
+from modyn.models.modyn_model import ModynModel
 from torch import nn
 from torchvision.models import densenet121
 
@@ -19,7 +20,7 @@ def __init__(self, model_configuration: dict[str, Any], device: str, amp: bool)
         self.model.to(device)
 
 
-class FmowNetModel(nn.Module):
+class FmowNetModel(ModynModel):
     def __init__(self, num_classes: int) -> None:
         super().__init__()
         self.num_classes = num_classes
@@ -31,5 +32,8 @@ def forward(self, data: torch.Tensor) -> torch.Tensor:
         out = F.relu(features, inplace=True)
         out = F.adaptive_avg_pool2d(out, (1, 1))
         out = torch.flatten(out, 1)
-
+        out = self.embedding_recorder(out)
         return self.classifier(out)
+
+    def get_last_layer(self) -> nn.Module:
+        return self.classifier

modyn/models/modyn_model.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+import torch
+from torch import nn
+
+
+# acknowledgment: github.com/PatrickZH/DeepCore/
+class EmbeddingRecorder(nn.Module):
+    def __init__(self, record_embedding: bool = False):
+        super().__init__()
+        self.record_embedding = record_embedding
+        self.embedding: Optional[torch.Tensor] = None
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        if self.record_embedding:
+            self.embedding = tensor
+        return tensor
+
+    def __enter__(self) -> None:
+        self.record_embedding = True
+
+    def __exit__(self, *args: Any) -> None:
+        self.record_embedding = False
+        self.embedding = None
+
+
+class ModynModel(nn.Module, ABC):
+    def __init__(self, record_embedding: bool = False) -> None:
+        super().__init__()
+        self.embedding_recorder = EmbeddingRecorder(record_embedding)
+
+    @property
+    def embedding(self) -> Optional[torch.Tensor]:
+        assert self.embedding_recorder is not None
+        return self.embedding_recorder.embedding
+
+    @abstractmethod
+    def get_last_layer(self) -> nn.Module:
+        raise NotImplementedError()

modyn/models/resnet18/resnet18.py

@@ -1,10 +1,44 @@
 from typing import Any
 
-from torchvision import models
+import torch
+from modyn.models.modyn_model import ModynModel
+from torch import Tensor, nn
+from torchvision.models.resnet import BasicBlock, ResNet
 
 
 class ResNet18:
     # pylint: disable-next=unused-argument
     def __init__(self, model_configuration: dict[str, Any], device: str, amp: bool) -> None:
-        self.model = models.__dict__["resnet18"](**model_configuration)
+        self.model = ResNetModyn(model_configuration)
         self.model.to(device)
+
+
+# the following class is adapted from
+# torchvision https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py
+
+
+class ResNetModyn(ResNet, ModynModel):
+    def __init__(self, model_configuration: dict[str, Any]) -> None:
+        super().__init__(BasicBlock, [2, 2, 2, 2], **model_configuration)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        # the following line is the only difference compared to the original implementation
+        x = self.embedding_recorder(x)
+        x = self.fc(x)
+
+        return x
+
+    def get_last_layer(self) -> nn.Module:
+        return self.fc

modyn/models/yearbooknet/yearbooknet.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 import torch
+from modyn.models.modyn_model import ModynModel
 from torch import nn
 
 
@@ -17,7 +18,7 @@ def __init__(self, model_configuration: dict[str, Any], device: str, amp: bool)
         self.model.to(device)
 
 
-class YearbookNetModel(nn.Module):
+class YearbookNetModel(ModynModel):
     def __init__(self, num_input_channels: int, num_classes: int) -> None:
         super().__init__()
         self.enc = nn.Sequential(
@@ -37,5 +38,8 @@ def conv_block(self, in_channels: int, out_channels: int) -> nn.Module:
     def forward(self, data: torch.Tensor) -> torch.Tensor:
         data = self.enc(data)
         data = torch.mean(data, dim=(2, 3))
-
+        data = self.embedding_recorder(data)
         return self.classifier(data)
+
+    def get_last_layer(self) -> nn.Module:
+        return self.classifier

modyn/tests/models/test_dlrm.py

@@ -67,3 +67,14 @@ def test_dlrm_reorder_categorical_input():
     assert reordered_test_data.shape == (64, 26)
     assert reordered_test_data.dtype == torch.long
     assert torch.equal(reordered_test_data, input_data)
+
+
+def test_get_last_layer():
+    net = DLRM(get_dlrm_configuration(), "cpu", False)
+    last_layer = net.model.get_last_layer()
+
+    assert isinstance(last_layer, torch.nn.Linear)
+    assert last_layer.in_features == 16
+    assert last_layer.out_features == 1
+    assert last_layer.bias.shape == (1,)
+    assert last_layer.weight.shape == (1, 16)

modyn/tests/models/test_embedding_recorder.py

@@ -0,0 +1,33 @@
+import torch
+from modyn.models.modyn_model import EmbeddingRecorder
+
+
+def test_embedding_recording():
+    recorder = EmbeddingRecorder()
+    with recorder:
+        input_tensor = torch.tensor([1, 2, 3])
+        output_tensor = recorder(input_tensor)
+        assert torch.equal(recorder.embedding, input_tensor)
+        assert torch.equal(output_tensor, input_tensor)
+
+
+def test_no_embedding_recording():
+    recorder = EmbeddingRecorder()
+    input_tensor = torch.tensor([4, 5, 6])
+    output_tensor = recorder(input_tensor)
+    assert recorder.embedding is None
+    assert torch.equal(output_tensor, input_tensor)
+
+
+def test_toggle_embedding_recording():
+    recorder = EmbeddingRecorder()
+    with recorder:
+        input_tensor = torch.tensor([7, 8, 9])
+        output_tensor = recorder(input_tensor)
+        assert torch.equal(recorder.embedding, input_tensor)
+        assert torch.equal(output_tensor, input_tensor)
+
+    input_tensor = torch.tensor([10, 11, 12])
+    output_tensor = recorder(input_tensor)
+    assert recorder.embedding is None
+    assert torch.equal(output_tensor, input_tensor)

modyn/tests/models/test_yearbook_net.py

@@ -42,3 +42,34 @@ def test_model_conv_block():
 
     # Assert that the output has the correct shape
     assert output.shape == (batch_size, 32, height // 2, width // 2)
+
+
+def test_forward_with_embedding_recording():
+    net = YearbookNet({"num_input_channels": 3, "num_classes": 10}, "cpu", False)
+    input_data = torch.rand(30, 3, 32, 32)
+    output = net.model(input_data)
+    assert output.shape == (30, 10)
+    assert net.model.embedding is None
+
+    with net.model.embedding_recorder:
+        input_data = torch.rand(30, 3, 32, 32)
+        output = net.model(input_data)
+        assert output.shape == (30, 10)
+
+        # expected embedding
+        expected_embedding = net.model.enc(input_data)
+        expected_embedding = torch.mean(expected_embedding, dim=(2, 3))
+
+        assert net.model.embedding is not None
+        assert torch.equal(net.model.embedding, expected_embedding)
+
+
+def test_get_last_layer():
+    net = YearbookNet({"num_input_channels": 3, "num_classes": 10}, "cpu", False)
+    last_layer = net.model.get_last_layer()
+
+    assert isinstance(last_layer, torch.nn.Linear)
+    assert last_layer.in_features == 32
+    assert last_layer.out_features == 10
+    assert last_layer.bias.shape == (10,)
+    assert last_layer.weight.shape == (10, 32)