fix SDE class bugs in utils and mnist notebooks (#87)

Fixes SDE class bugs in utils and mnist notebooks. Refactors the example folder to split the examples into different folders. Adds two new tutorials (Flow Matching and minibatch OT)
atong01 · Dec 15, 2023 · 6b3adb4 · 6b3adb4
1 parent b29d418
commit 6b3adb4
Show file tree

Hide file tree

Showing 22 changed files with 3,205 additions and 1,892 deletions.
diff --git a/examples/2D_tutorials/Flow_matching_tutorial.ipynb b/examples/2D_tutorials/Flow_matching_tutorial.ipynb
diff --git a/examples/2D_tutorials/SF2M_tutorial.ipynb b/examples/2D_tutorials/SF2M_tutorial.ipynb
diff --git a/examples/2D_tutorials/The_unreasonable_performance_of_minibatch_OT.ipynb b/examples/2D_tutorials/The_unreasonable_performance_of_minibatch_OT.ipynb
diff --git a/...notebooks/model-comparison-plotting.ipynb → ...tutorials/model-comparison-plotting.ipynb b/...notebooks/model-comparison-plotting.ipynb → ...tutorials/model-comparison-plotting.ipynb
diff --git a/...sing/1.0-embryoid_body_data_to_h5ad.ipynb → ...sing/1.0-embryoid_body_data_to_h5ad.ipynb b/...sing/1.0-embryoid_body_data_to_h5ad.ipynb → ...sing/1.0-embryoid_body_data_to_h5ad.ipynb
diff --git a/.../preprocessing/1.1-eb-data-1d-phate.ipynb → .../preprocessing/1.1-eb-data-1d-phate.ipynb b/.../preprocessing/1.1-eb-data-1d-phate.ipynb → .../preprocessing/1.1-eb-data-1d-phate.ipynb
diff --git a/...s/preprocessing/2.1-split-cite-data.ipynb → ...s/preprocessing/2.1-split-cite-data.ipynb b/...s/preprocessing/2.1-split-cite-data.ipynb → ...s/preprocessing/2.1-split-cite-data.ipynb
diff --git a/examples/notebooks/preprocessing/README.md → ...ples/2D_tutorials/preprocessing/README.md b/examples/notebooks/preprocessing/README.md → ...ples/2D_tutorials/preprocessing/README.md
diff --git a/...ebooks/training-8gaussians-to-moons.ipynb → ...otial_training_8_gaussians_to_moons.ipynb b/...ebooks/training-8gaussians-to-moons.ipynb → ...otial_training_8_gaussians_to_moons.ipynb
diff --git a/examples/cifar10/README.md → examples/images/cifar10/README.md b/examples/cifar10/README.md → examples/images/cifar10/README.md
@@ -3,7 +3,7 @@
 This repository is used to reproduce the CIFAR-10 experiments from [1](https://arxiv.org/abs/2302.00482). We have designed a novel experimental procedure that helps us to reach an __FID of 3.5__ on the Cifar10 dataset.
 
 <p align="center">
-<img src="../../assets/169_generated_samples_otcfm.png" width="600"/>
+<img src="../../../assets/169_generated_samples_otcfm.png" width="600"/>
 </p>
 
 To reproduce the experiments and save the weights, install the requirements from the main repository and then run (runs on a single RTX 2080 GPU):

diff --git a/examples/cifar10/compute_fid.py → examples/images/cifar10/compute_fid.py b/examples/cifar10/compute_fid.py → examples/images/cifar10/compute_fid.py
diff --git a/examples/cifar10/train_cifar10.py → examples/images/cifar10/train_cifar10.py b/examples/cifar10/train_cifar10.py → examples/images/cifar10/train_cifar10.py
diff --git a/examples/cifar10/utils_cifar.py → examples/images/cifar10/utils_cifar.py b/examples/cifar10/utils_cifar.py → examples/images/cifar10/utils_cifar.py
diff --git a/examples/images/conditional_mnist.ipynb b/examples/images/conditional_mnist.ipynb
diff --git a/examples/images/mnist_example.ipynb b/examples/images/mnist_example.ipynb
diff --git a/examples/notebooks/SF2M_2D_example.ipynb b/examples/notebooks/SF2M_2D_example.ipynb
diff --git a/examples/notebooks/conditional_mnist.ipynb b/examples/notebooks/conditional_mnist.ipynb
diff --git a/examples/notebooks/mnist_example.ipynb b/examples/notebooks/mnist_example.ipynb
diff --git a/examples/notebooks/single-cell_example.ipynb b/examples/notebooks/single-cell_example.ipynb
diff --git a/examples/single_cell/single-cell_example.ipynb b/examples/single_cell/single-cell_example.ipynb
diff --git a/torchcfm/utils.py b/torchcfm/utils.py
@@ -64,31 +64,3 @@ def plot_trajectories(traj):
     plt.xticks([])
     plt.yticks([])
     plt.show()
-
-
-class SDE(torch.nn.Module):
-    noise_type = "diagonal"
-    sde_type = "ito"
-
-    def __init__(self, ode_drift, score, noise=1.0, reverse=False):
-        super().__init__()
-        self.drift = ode_drift
-        self.score = score
-        self.reverse = reverse
-        self.noise = noise
-
-    # Drift
-    def f(self, t, y):
-        if self.reverse:
-            t = 1 - t
-        if len(t.shape) == len(y.shape):
-            x = torch.cat([y, t], 1)
-        else:
-            x = torch.cat([y, t.repeat(y.shape[0])[:, None]], 1)
-        if self.reverse:
-            return -self.drift(x) + self.score(x)
-        return self.drift(x) + self.score(x)
-
-    # Diffusion
-    def g(self, t, y):
-        return torch.ones_like(t) * torch.ones_like(y) * self.noise
diff --git a/torchcfm/version.py b/torchcfm/version.py
@@ -1 +1 @@
-__version__ = "1.0.5"
+__version__ = "1.0.6"