DiT, a bit of clean up

Author: homerjed

.gitignore

@@ -4,6 +4,7 @@ __pycache__/
 imgs/
 exps/
 _fisher.py
+sbgm/_sbgm.py
 _set_transformer.py
 .pytest_cacche/
 __unet.py

data/cifar10.py

@@ -18,21 +18,23 @@ def convert_torch_to_in_memory(dataset):
 
 
 def cifar10(path: str, key: Key, *, in_memory: bool = True) -> ScalerDataset:
+
     key_train, key_valid = jr.split(key)
 
     n_pix = 32 # Native resolution for CIFAR10 
     data_shape = (3, n_pix, n_pix)
+    context_shape = None
     parameter_dim = 1
     n_classes = 10
 
-    scaler = Scaler(x_min=0., x_max=1.)
+    scaler = Normer()
 
     train_transform = transforms.Compose(
         [
             transforms.Resize((n_pix, n_pix)),
             transforms.RandomHorizontalFlip(),
             transforms.ToTensor(), 
-            transforms.Lambda(scaler.forward) # [0,1] -> [-1,1]
+            transforms.Lambda(scaler.forward) 
         ]
     )
     valid_transform = transforms.Compose(
@@ -64,7 +66,6 @@ def cifar10(path: str, key: Key, *, in_memory: bool = True) -> ScalerDataset:
         At = At.astype(jnp.float32)
         Av = Av.astype(jnp.float32)
 
-        # process_fn = Scaler(x_min=Xt.min(), x_max=Xt.max())
         process_fn = Normer(x_mean=Xt.mean(), x_std=Xt.std())
 
         train_dataloader = InMemoryDataLoader(
@@ -76,10 +77,18 @@ def cifar10(path: str, key: Key, *, in_memory: bool = True) -> ScalerDataset:
         process_fn = Scaler(x_min=0., x_max=1.)
 
         train_dataloader = TorchDataLoader(
-            train_dataset, data_shape, parameter_dim=parameter_dim, key=key_train
+            train_dataset, 
+            data_shape=data_shape, 
+            context_shape=context_shape,
+            parameter_dim=parameter_dim, 
+            key=key_train
         )
         valid_dataloader = TorchDataLoader(
-            valid_dataset, data_shape, parameter_dim=parameter_dim, key=key_valid
+            valid_dataset, 
+            data_shape=data_shape, 
+            context_shape=context_shape,
+            parameter_dim=parameter_dim, 
+            key=key_valid
         )
 
     def label_fn(key, n):

data/flowers.py

@@ -8,7 +8,9 @@
 
 
 def flowers(path: str, key: Key, n_pix: int) -> ScalerDataset:
+
     key_train, key_valid = jr.split(key)
+
     data_shape = (3, n_pix, n_pix)
     parameter_dim = 1
     n_classes = 102
@@ -19,7 +21,6 @@ def flowers(path: str, key: Key, n_pix: int) -> ScalerDataset:
         [
             transforms.Resize((n_pix, n_pix)),
             transforms.RandomCrop(n_pix, padding=4, padding_mode='reflect'),
-            # transforms.Grayscale(),
             transforms.RandomHorizontalFlip(),
             transforms.RandomVerticalFlip(),
             transforms.ToTensor(),
@@ -30,11 +31,11 @@ def flowers(path: str, key: Key, n_pix: int) -> ScalerDataset:
         [
             transforms.Resize((n_pix, n_pix)),
             transforms.RandomCrop(n_pix, padding=4, padding_mode='reflect'),
-            # transforms.Grayscale(),
             transforms.ToTensor(),
             transforms.Lambda(scaler.forward)
         ]
     )
+
     train_dataset = datasets.Flowers102(
         os.path.join(path, "datasets/flowers/"), 
         split="train", 

data/grfs.py

@@ -9,7 +9,7 @@
 from torchvision import transforms
 import powerbox
 
-from .utils import Scaler, ScalerDataset, TorchDataLoader, InMemoryDataLoader
+from .utils import Scaler, Normer, ScalerDataset, TorchDataLoader, InMemoryDataLoader
 
 data_dir = "/project/ls-gruen/users/jed.homer/data/fields/"
 
@@ -114,15 +114,10 @@ def grfs(
 
     print("\nFields data:", X.shape, Q.shape)
 
-    min = X.min()
-    max = X.max()
-    X = (X - min) / (max - min) # ... -> [0, 1]
+    X = (X - jnp.mean(X, axis=0)) / jnp.std(X, axis=0)  # Standardize fields
+    Q = (Q - jnp.mean(Q, axis=0)) / jnp.std(Q, axis=0)  # Standardize fields
 
-    # min = Q.min()
-    # max = Q.max()
-    # Q = (Q - min) / (max - min) # ... -> [0, 1]
-
-    scaler = Scaler() # [0,1] -> [-1,1]
+    scaler = Normer() #Scaler() # [0,1] -> [-1,1]
 
     n_train = int(split * n_fields)
 
@@ -152,18 +147,24 @@ def grfs(
             (X[n_train:], Q[n_train:], A[n_train:]), transform=valid_transform
         )
         train_dataloader = TorchDataLoader(
-            train_dataset, data_shape, parameter_dim=parameter_dim, key=key_train
+            train_dataset, 
+            data_shape=data_shape, 
+            context_shape=context_shape, 
+            parameter_dim=parameter_dim, 
+            key=key_train
         )
         valid_dataloader = TorchDataLoader(
-            valid_dataset, data_shape, parameter_dim=parameter_dim, key=key_valid
+            valid_dataset, 
+            data_shape=data_shape, 
+            context_shape=context_shape,
+            parameter_dim=parameter_dim, 
+            key=key_valid
         )
 
     def label_fn(key: Key[jnp.ndarray, "..."], n: int) -> Tuple[Array, Array]:
         Q, A = get_grf_labels(n_pix)
         ix = jr.choice(key, jnp.arange(len(Q)), (n,))
-        Q = Q[ix]
-        A = A[ix]
-        return Q, A
+        return Q[ix], A[ix]
 
     return ScalerDataset(
         name="grfs",

data/mnist.py

@@ -13,7 +13,9 @@ def tensor_to_array(tensor: Tensor) -> Array:
 
 
 def mnist(path:str, key: Key, *, in_memory: bool = True) -> ScalerDataset:
+
     key_train, key_valid = jr.split(key)
+
     n_pix = 28
     data_shape = (1, n_pix, n_pix)
     parameter_dim = 1 

data/moons.py

@@ -10,6 +10,7 @@ def key_to_seed(key):
 
 def moons(key):
     key_train, key_valid = jr.split(key)
+
     data_shape = (2,)
     context_shape = None 
     parameter_dim = 1
@@ -24,17 +25,6 @@ def moons(key):
     Yv = Yv[:, jnp.newaxis].astype(jnp.float32)
 
     process_fn = Normer(Xt.mean(), Xt.std())
-
-    # min = Xt.min()
-    # max = Xt.max()
-    # mean = Xt.mean()
-    # std = Xt.std()
-
-    # (We do need to handle normalisation ourselves though.)
-    # train_data = (Xt - min) / (max - min)
-    # valid_data = (Xv - min) / (max - min)
-    # train_data = (Xt - mean) / std
-    # valid_data = (Xv - mean) / std
 
     train_dataloader = InMemoryDataLoader(
         X=jnp.asarray(Xt), Q=jnp.asarray(Yt), A=None, process_fn=process_fn, key=key_train