Start data loading and training process

scripts/train.py

@@ -1,7 +1,23 @@
+from torch.utils.data import DataLoader
+import torch
+from accelerate import Accelerator
+from tqdm import tqdm
+from datasets import load_dataset
 from umi.config import Config
 from umi.models.unet import create_model
+from umi.datasets import CIFAR10Dataset
+from diffusers import DDPMPipeline, DDPMScheduler
 
 if __name__ == '__main__':
-    model_config = Config()
+    config = Config()
 
-    model = create_model(model_config)
+    model = create_model(config)
+    dataset = load_dataset("cifar10", split="train")
+    dataset = CIFAR10Dataset(dataset, transform=config.transform)
+    dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
+
+    noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+    accelerator = Accelerator()
+    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

umi/config.py

@@ -1,6 +1,8 @@
 from torchvision import transforms
 
 class Config:
+    epochs :int = 4
+
     transform = transforms.Compose([
         transforms.Resize(32),
         transforms.CenterCrop(32),

umi/datasets/__init__.py

@@ -0,0 +1 @@
+from cifar10 import CIFAR10Dataset

umi/datasets/cifar10.py

@@ -0,0 +1,15 @@
+from torch.utils.data import Dataset
+
+class CIFAR10Dataset(Dataset):
+    def __init__(self, dataset, transform=None):
+        self.dataset = dataset
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        img = self.dataset[idx]["img"]
+        if self.transform:
+            img = self.transform(img)
+        return {"image": img}