Add boilerplate code for Dataflux-Pytorch Lightning demo

demo/lightning/data.py

@@ -0,0 +1,24 @@
+import torch
+import lightning.pytorch as pl
+
+from dataflux_pytorch import dataflux_mapstyle_dataset
+
+
+class Unet3DDataModule(pl.LighitningDataModule):
+    def __init__(self, gcs_bucket):
+        self.data_dir = gcs_bucket + "/images"
+        self.transform = 1
+
+    def prepare_data(self):
+
+        pass
+
+    def setup(self, state="train"):
+        # Init DatafluxPytTrain object 
+        # Pass transform function as argument - modify DatafluxPyTrain constructor
+        pass
+
+    def train_dataloader(self):
+        # Init dataflux dataloader
+        # Wrap it in torch.DataLoader
+        pass

demo/lightning/model.py

@@ -0,0 +1,22 @@
+import torch
+import lightning.pytorch as pl
+
+
+class Unet3D(pl.LightningModule):
+    def __init__(self):
+        pass
+
+    def forward(self, x):
+        pass
+
+    def configure_optimizers(self):
+        pass
+
+    def training_step(self, train_batch, batch_idx):
+        pass
+
+    def backward(self, trainer, loss, optimizer, optimizer_idx):
+        pass
+
+    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx):
+        pass

demo/lightning/train.py

@@ -0,0 +1,14 @@
+import lightning as pl
+
+from model import Unet3D
+from data import Unet3DDataModule
+
+if __name__ == "__main__":
+    model = Unet3D()
+    train_data_loader = Unet3DDataModule()
+    trainer = pl.Trainer(
+        devices=2, 
+        accelerator="gpu", 
+        max_epochs=5
+    )
+    trainer.fit(model=model, train_dataloaders=train_data_loader)