Tests on checkerboard

Author: Akanota

configs/experiments/checkerboard_steps/flow_matching.yaml

@@ -0,0 +1,85 @@
+name: checkerboard-fm
+project: flowmaps
+
+dim_flow: 2
+dim_conditioning: 0
+dim_network: 2
+batch_size: 256
+
+weighting: 'one'
+loss_type: 'fm_ot'
+scaling_type: 'none'
+
+learning_rate: 0.001
+
+hidden_dims: [32, 64, 128, 256, 128, 64, 32]
+
+data:
+  target: sbisim.data.benchmarks.density_estimation.ToyDataset
+  _file: configs/datasets/2d/checkerboard.yaml
+
+strategy:
+  target: sbisim.strategy.ConditionalFlowMatching
+  params:
+    dim_flow: ${dim_flow}
+    dim_conditioning: ${dim_conditioning}
+    weighting: ${weighting}
+    loss_type: ${loss_type}
+    scaling_type: ${scaling_type}
+    time_alpha: 4
+    bandwidth: 0.01
+    prior:
+      target: sbisim.strategy.distributions.get_normal_gaussian_prior
+      params:
+        shape: [ 2 ]
+    sampler:
+      target: sbisim.strategy.improved_inference.PriorODESolver
+      params:
+        prior:
+          target: sbisim.strategy.distributions.get_normal_gaussian_prior
+          params:
+            shape: [ 2 ]
+        solver_name: euler
+        init_stepsize: 0.01
+    schedule:
+        target: sbisim.strategy.get_paths
+        params:
+          paths:
+            name: optimal_transport
+            params:
+              sigma_min: 1e-4
+    model:
+      target: sbisim.flows.models.fmpe.SBIResidualNet
+      params:
+        hidden_dims: ${hidden_dims}
+        out_dim: ${dim_flow}
+        in_dim: ${dim_network}
+        activation_fn: gelu
+
+num_steps: 10000
+training:
+  num_epochs: 15
+  num_steps_per_epoch: ${num_steps}
+  batch_size: ${batch_size}
+  early_stopping:
+    patience: 50
+
+test:
+  active: True
+  tests:
+    scatter:
+      target: sbisim.callbacks.TwoDimensionalPlot
+      params:
+        savedir: ${runtime.logdir}
+
+
+patience: 3
+optimization:
+  _file: configs/optimization/adam_reduce_lr.yaml
+
+callbacks:
+  checkpoint:
+    target: sbisim.callbacks.checkpoint.Checkpoint
+    params:
+      savedir: ${runtime.logdir}
+      key: val_loss

requirements.txt

@@ -14,4 +14,5 @@ objax
 corner
 orbax
 pandas
-astropy
+astropy
+seaborn

sbisim/callbacks/__init__.py

@@ -1,2 +1,3 @@
 from .c2st import C2ST
-from .visualization import BenchmarkScatterPlot
+from .visualization import BenchmarkScatterPlot
+from .two_dimensional import TwoDimensionalPlot

sbisim/callbacks/two_dimensional.py

@@ -9,28 +9,30 @@
 
 from ..data.density_estimation.golden_ratio import fibonacci_ratio, get_centers
 
-# from source.data.dataloader import DataLoader
 from ..strategy import Strategy
-
 from matplotlib import pyplot as plt
-
 import jax.numpy as jnp
 from tqdm import tqdm
-
 import seaborn as sns
 
+
+class DataLoader:
+    pass
+
+
 class TwoDimensionalPlot(Callback):
 
     name: str = '2d_plot'
     save_every: int = 10
     num_samples: int = 3000
     num_samples_gt: int = 3000
 
-    def __init__(self, save_every: int = 10, num_samples: int = 3000):
+    def __init__(self, save_every: int = 10, num_samples: int = 3000, savedir: str = None):
         super(self.__class__, self).__init__()
 
         self.save_every = save_every
         self.num_samples = num_samples
+        self.savedir = savedir
 
 
     def __call__(self, logs: dict, rng: jr.PRNGKey, *args, **kwargs):
@@ -59,8 +61,12 @@ def _call(self, logs: dict, rng: jr.PRNGKey, strategy: Strategy,
 
             ax.scatter(x[:self.num_samples_gt], y[:self.num_samples_gt])
 
+            if self.savedir is not None:
+                plt.savefig(f"{self.savedir}/pictures/samples_gt.png")
+
             logs[f'samples_gt'] = wandb.Image(fig)
 
+
             plt.close(fig)
 
         sample_fn = strategy.sample
@@ -75,6 +81,15 @@ def _call(self, logs: dict, rng: jr.PRNGKey, strategy: Strategy,
         fig, ax = plt.subplots()
         ax.scatter(x, y)
 
+        epoch = logs.get("epoch", 0)
+        if self.savedir is not None:
+
+            import os
+            # create directory if it does not exist
+            os.makedirs(f"{self.savedir}/pictures", exist_ok=True)
+
+            plt.savefig(f"{self.savedir}/pictures/samples_{epoch}.png")
+
         logs[f'samples'] = wandb.Image(fig)
 
         plt.close(fig)
@@ -84,6 +99,9 @@ def _call(self, logs: dict, rng: jr.PRNGKey, strategy: Strategy,
     def on_train_begin(self, *args, **kwargs):
         return self.__call__(*args, init=True, **kwargs)
 
+    def on_test(self, *args, **kwargs):
+        return self.__call__(*args, init=False, **kwargs)
+
     def on_epoch_end(self, logs: dict, rng: jr.PRNGKey, *args, **kwargs):
 
         if "epoch" in logs and logs["epoch"] % self.save_every == 0:

sbisim/trainer.py

@@ -84,6 +84,14 @@ def init_strategy(self):
 
         print("Number of parameters: ", wandb.run.summary["num_params"])
 
+        # Check if JAX is using GPU
+        from jax.lib import xla_bridge
+        platform = xla_bridge.get_backend().platform
+        if platform == 'gpu':
+            print("JAX is using GPU")
+        else:
+            print(f"JAX is using {platform}")
+
     def pack_(self):
         self.strategy.bind(self.opt.get_params())
         return {'rng': self.callback_rng, 'strategy': self.strategy,
@@ -212,14 +220,6 @@ def train_epoch(self, logs, rng):
                 step_ += 1
                 if step_ >= num_val_batches:
                     break
-                
-            # Check if JAX is using GPU
-            from jax.lib import xla_bridge
-            platform = xla_bridge.get_backend().platform
-            if platform == 'gpu':
-                print("JAX is using GPU")
-            else:
-                print(f"JAX is using {platform}")
 
             avg_loss /= step_