[test] refactored testing components (#324)

tests/components_to_test/__init__.py

@@ -0,0 +1 @@
+from . import repeated_computed_layer, resnet, nested_model

tests/components_to_test/nested_model.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .utils import DummyDataGenerator
+from .registry import non_distributed_component_funcs
+
+
+class SubNet(nn.Module):
+
+    def __init__(self, out_features) -> None:
+        super().__init__()
+        self.bias = nn.Parameter(torch.zeros(out_features))
+
+    def forward(self, x, weight):
+        return F.linear(x, weight, self.bias)
+
+
+class NestedNet(nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(5, 5)
+        self.sub_fc = SubNet(5)
+        self.fc2 = nn.Linear(5, 2)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.sub_fc(x, self.fc1.weight)
+        x = self.fc1(x)
+        x = self.fc2(x)
+        return x
+
+
+class DummyDataLoader(DummyDataGenerator):
+
+    def generate(self):
+        data = torch.rand(16, 5)
+        label = torch.randint(low=0, high=2, size=(16,))
+        return data, label
+
+
+@non_distributed_component_funcs.register(name='nested_model')
+def get_training_components():
+    model = NestedNet()
+    trainloader = DummyDataLoader()
+    testloader = DummyDataLoader()
+    optim = torch.optim.Adam(model.parameters(), lr=0.001)
+    criterion = torch.nn.CrossEntropyLoss()
+    return model, trainloader, testloader, optim, criterion

tests/components_to_test/registry.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+
+
+class Registry:
+
+    def __init__(self):
+        self._registry = dict()
+
+    def register(self, name):
+        assert name not in self._registry
+
+        def _regsiter(callable_):
+            self._registry[name] = callable_
+
+        return _regsiter
+
+    def get_callable(self, name: str):
+        return self._registry[name]
+
+    def __iter__(self):
+        self._idx = 0
+        self._len = len(self._registry)
+        self._names = list(self._registry.keys())
+        return self
+
+    def __next__(self):
+        if self._idx < self._len:
+            key = self._names[self._idx]
+            callable_ = self._registry[key]
+            self._idx += 1
+            return callable_
+        else:
+            raise StopIteration
+
+
+non_distributed_component_funcs = Registry()
+model_paralle_component_funcs = Registry()
+
+__all__ = ['non_distributed_component_funcs', 'model_paralle_component_funcs']

tests/components_to_test/repeated_computed_layer.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python
+
+import torch
+import torch.nn as nn
+from colossalai.nn import CheckpointModule
+from .utils.dummy_data_generator import DummyDataGenerator
+from .registry import non_distributed_component_funcs
+
+
+class NetWithRepeatedlyComputedLayers(CheckpointModule):
+    """
+    This model is to test with layers which go through forward pass multiple times.
+    In this model, the fc1 and fc2 call forward twice
+    """
+
+    def __init__(self, checkpoint=False) -> None:
+        super().__init__(checkpoint=checkpoint)
+        self.fc1 = nn.Linear(5, 5)
+        self.fc2 = nn.Linear(5, 5)
+        self.fc3 = nn.Linear(5, 2)
+        self.layers = [self.fc1, self.fc2, self.fc1, self.fc2, self.fc3]
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class DummyDataLoader(DummyDataGenerator):
+
+    def generate(self):
+        data = torch.rand(16, 5)
+        label = torch.randint(low=0, high=2, size=(16,))
+        return data, label
+
+
+@non_distributed_component_funcs.register(name='repeated_computed_layers')
+def get_training_components():
+    model = NetWithRepeatedlyComputedLayers(checkpoint=True)
+    trainloader = DummyDataLoader()
+    testloader = DummyDataLoader()
+    optim = torch.optim.Adam(model.parameters(), lr=0.001)
+    criterion = torch.nn.CrossEntropyLoss()
+    return model, trainloader, testloader, optim, criterion

tests/components_to_test/resnet.py

@@ -0,0 +1,30 @@
+from torchvision.models import resnet18
+from .registry import non_distributed_component_funcs
+from pathlib import Path
+import os
+import torch
+from torchvision.transforms import transforms
+from torchvision.datasets import CIFAR10
+from colossalai.utils import get_dataloader
+
+
+def get_cifar10_dataloader(train):
+    # build dataloaders
+    dataset = CIFAR10(root=Path(os.environ['DATA']),
+                      download=True,
+                      train=train,
+                      transform=transforms.Compose(
+                          [transforms.ToTensor(),
+                           transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))]))
+    dataloader = get_dataloader(dataset=dataset, shuffle=True, batch_size=16, drop_last=True)
+    return dataloader
+
+
+@non_distributed_component_funcs.register(name='resnet18')
+def get_resnet_training_components():
+    model = resnet18(num_classes=10)
+    trainloader = get_cifar10_dataloader(train=True)
+    testloader = get_cifar10_dataloader(train=False)
+    optim = torch.optim.Adam(model.parameters(), lr=0.001)
+    criterion = torch.nn.CrossEntropyLoss()
+    return model, trainloader, testloader, optim, criterion

tests/components_to_test/utils/__init__.py

@@ -0,0 +1 @@
+from .dummy_data_generator import DummyDataGenerator

tests/components_to_test/utils/dummy_data_generator.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+
+
+class DummyDataGenerator(ABC):
+
+    @abstractmethod
+    def generate(self):
+        pass
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return self.generate()