add stochastic interpolants

atong01 · Jan 16, 2024 · 952bdec · 952bdec
1 parent cfafc28
commit 952bdec
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/examples/images/cifar10/train_cifar10.py b/examples/images/cifar10/train_cifar10.py
@@ -17,6 +17,7 @@
     ConditionalFlowMatcher,
     ExactOptimalTransportConditionalFlowMatcher,
     TargetConditionalFlowMatcher,
+    VariancePreservingConditionalFlowMatcher,
 )
 from torchcfm.models.unet.unet import UNetModelWrapper
 
@@ -128,9 +129,11 @@ def train(argv):
         FM = ConditionalFlowMatcher(sigma=sigma)
     elif FLAGS.model == "fm":
         FM = TargetConditionalFlowMatcher(sigma=sigma)
+    elif FLAGS.model == "si":
+        FM = VariancePreservingConditionalFlowMatcher(sigma=sigma)
     else:
         raise NotImplementedError(
-            f"Unknown model {FLAGS.model}, must be one of ['otcfm', 'icfm', 'fm']"
+            f"Unknown model {FLAGS.model}, must be one of ['otcfm', 'icfm', 'fm', 'si']"
         )
 
     savedir = FLAGS.output_dir + FLAGS.model + "/"
@@ -162,7 +165,7 @@ def train(argv):
                         "optim": optim.state_dict(),
                         "step": step,
                     },
-                    savedir + f"cifar10_weights_step_{step}.pt",
+                    savedir + f"{FLAGS.model}_cifar10_weights_step_{step}.pt",
                 )