refactor cifar10 example

client/starwhale/core/dataset/type.py

@@ -238,6 +238,7 @@ def astype(self) -> t.Dict[str, t.Any]:
             "mime_type": self.mime_type,
             "shape": self.shape,
             "encoding": self.encoding,
+            "display_name": self.display_name,
         }
 
     def asdict(self, ignore_keys: t.Optional[t.List[str]] = None) -> t.Dict[str, t.Any]:

example/cifar10/.gitignore

@@ -0,0 +1,2 @@
+data/
+models/

example/cifar10/Makefile

@@ -0,0 +1,5 @@
+.POHNY: train
+train:
+	mkdir -p models data
+	python3 cifar/train.py
+

example/cifar10/cifar/dataset.py

@@ -0,0 +1,54 @@
+import io
+import pickle
+import typing as t
+from pathlib import Path
+
+from PIL import Image as PILImage
+
+from starwhale.api.dataset import Image, MIMEType, BuildExecutor
+
+ROOT_DIR = Path(__file__).parent.parent / "data" / "cifar-10-batches-py"
+TRAIN_DATASET_PATHS = [ROOT_DIR / f"data_batch_{i}" for i in range(1, 6)]
+TEST_DATASET_PATHS = [ROOT_DIR / "test_batch"]
+
+
+def parse_meta() -> t.Dict[str, t.Any]:
+    with (ROOT_DIR / "batches.meta").open("rb") as f:
+        return pickle.load(f)
+
+
+dataset_meta = parse_meta()
+
+
+def _iter_item(paths: t.List[Path]) -> t.Generator[t.Tuple[t.Any, t.Dict], None, None]:
+    for path in paths:
+        with path.open("rb") as f:
+            content = pickle.load(f, encoding="bytes")
+            for data, label, filename in zip(
+                content[b"data"], content[b"labels"], content[b"filenames"]
+            ):
+                annotations = {
+                    "label": label,
+                    "label_display_name": dataset_meta["label_names"][label],
+                }
+
+                image_array = data.reshape(3, 32, 32).transpose(1, 2, 0)
+                image_bytes = io.BytesIO()
+                PILImage.fromarray(image_array).save(image_bytes, format="PNG")
+
+                yield Image(
+                    fp=image_bytes.getvalue(),
+                    display_name=filename.decode(),
+                    shape=image_array.shape,
+                    mime_type=MIMEType.PNG,
+                ), annotations
+
+
+class CIFAR10TrainBuildExecutor(BuildExecutor):
+    def iter_item(self) -> t.Generator[t.Tuple[t.Any, t.Any], None, None]:
+        return _iter_item(TRAIN_DATASET_PATHS)
+
+
+class CIFAR10TestBuildExecutor(BuildExecutor):
+    def iter_item(self) -> t.Generator[t.Tuple[t.Any, t.Any], None, None]:
+        return _iter_item(TEST_DATASET_PATHS)

example/cifar10/code/train.py → example/cifar10/cifar/train.py

@@ -1,33 +1,37 @@
+import os
+
 import torch
 import torch.nn as nn
+import torch.optim as optim
 import torchvision
 import torchvision.transforms as transforms
-import torch.optim as optim
 from model import Net
 from torch.optim.lr_scheduler import StepLR
-import os
 
 # https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
 _ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
 _MODEL_PATH = os.path.join(_ROOT_DIR, "../models/cifar_net.pth")
+_DATA_DIR = os.path.join(_ROOT_DIR, "data")
 
 
-def main():
+def train():
     transform = transforms.Compose(
-        [transforms.ToTensor(),
-         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+        [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    )
     batch_size = 4
-    trainset = torchvision.datasets.CIFAR10(root='../data', train=True,
-                                        download=True, transform=transform)
-    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
-                                          shuffle=True, num_workers=2)
+    train_set = torchvision.datasets.CIFAR10(
+        root=_DATA_DIR, train=True, download=True, transform=transform
+    )
+    train_loader = torch.utils.data.DataLoader(
+        train_set, batch_size=batch_size, shuffle=True, num_workers=2
+    )
     net = Net()
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
     scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
     for epoch in range(10):  # loop over the dataset multiple times
         running_loss = 0.0
-        for i, data in enumerate(trainloader, 0):
+        for i, data in enumerate(train_loader, 0):
             # get the inputs; data is a list of [inputs, labels]
             inputs, labels = data
             # zero the parameter gradients
@@ -37,17 +41,15 @@ def main():
             loss = criterion(outputs, labels)
             loss.backward()
             optimizer.step()
-            # print statistics
             running_loss += loss.item()
-            if i % 2000 == 1999:    # print every 2000 mini-batches
-                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
+            if i % 2000 == 1999:  # print every 2000 mini-batches
+                print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}")
                 running_loss = 0.0
         scheduler.step()
 
-    print('Finished Training')
-
+    print("Finished Training")
     torch.save(net.state_dict(), _MODEL_PATH)
 
 
 if __name__ == "__main__":
-    main()
+    train()

example/cifar10/dataset.yaml

@@ -1,12 +1,10 @@
-name: cifar10
-
-process: code.data_slicer:CIFAR10Slicer
-pip_req: requirements.txt
+#name: cifar10-train
+#process: cifar.dataset:CIFAR10TrainBuildExecutor
 
+name: cifar10-test
+process: cifar.dataset:CIFAR10TestBuildExecutor
 desc: CIFAR10 data and label test dataset
-tag:
-  - bin
 
 attr:
   alignment_size: 4k
-  volume_size: 2M
+  volume_size: 10M

example/cifar10/model.yaml

@@ -1,16 +1,9 @@
 version: 1.0
 name: cifar_net
-
 model:
   - models/cifar_net.pth
 
-config:
-  - config/hyperparam.json
-
 run:
-  ppl: code.ppl:CIFAR10Inference
+  ppl: cifar.ppl:CIFAR10Inference
 
 desc: cifar10 by pytorch
-
-tag:
-  - multi_classification