precommit hooks

Author: psteinb

examples/images/conditional_mnist_noninteger.ipynb

@@ -97,7 +97,11 @@
     "from torchvision.utils import make_grid\n",
     "from tqdm import tqdm\n",
     "\n",
-    "from torchcfm.conditional_flow_matching import ConditionalFlowMatcher, ExactOptimalTransportConditionalFlowMatcher, SchrodingerBridgeConditionalFlowMatcher\n",
+    "from torchcfm.conditional_flow_matching import (\n",
+    "    ConditionalFlowMatcher,\n",
+    "    ExactOptimalTransportConditionalFlowMatcher,\n",
+    "    SchrodingerBridgeConditionalFlowMatcher,\n",
+    ")\n",
     "from torchcfm.models.unet import UNetModel\n",
     "\n",
     "savedir = \"models/cond_mnist\"\n",
@@ -154,10 +158,11 @@
     "#    Float Conditional CFM\n",
     "#################################\n",
     "\n",
+    "\n",
     "class embed_condition(torch.nn.Module):\n",
-    "    \"\"\" simple network to embed the condition, other architectures can be used too \"\"\"\n",
+    "    \"\"\"simple network to embed the condition, other architectures can be used too\"\"\"\n",
     "\n",
-    "    def __init__(self, input_dim=1, target_dim = 128):\n",
+    "    def __init__(self, input_dim=1, target_dim=128):\n",
     "        super().__init__()\n",
     "\n",
     "        self.model = torch.nn.Sequential(\n",
@@ -171,7 +176,6 @@
     "        return self.model(label)\n",
     "\n",
     "\n",
-    "\n",
     "sigma = 0.0\n",
     "model = UNetModel(\n",
     "    dim=(1, 28, 28), num_channels=32, num_res_blocks=1, embedding_net=embed_condition\n",
@@ -11920,7 +11924,9 @@
     "    for i, data in enumerate(train_loader):\n",
     "        optimizer.zero_grad()\n",
     "        x1 = data[0].to(device)\n",
-    "        y = data[1].float().to(device).reshape((batch_size, 1)) / 2. #just to have a floating point label\n",
+    "        y = (\n",
+    "            data[1].float().to(device).reshape((batch_size, 1)) / 2.0\n",
+    "        )  # just to have a floating point label\n",
     "        x0 = torch.randn_like(x1)\n",
     "        t, xt, ut = FM.sample_location_and_conditional_flow(x0, x1)\n",
     "        vt = model(t, xt, y)\n",
@@ -11956,8 +11962,10 @@
    ],
    "source": [
     "USE_TORCH_DIFFEQ = True\n",
-    "ntest = 10*10\n",
-    "generated_class_list = torch.arange(10, device=device).repeat(10).reshape((ntest, 1)).float() / 2. #TODO: reshape\n",
+    "ntest = 10 * 10\n",
+    "generated_class_list = (\n",
+    "    torch.arange(10, device=device).repeat(10).reshape((ntest, 1)).float() / 2.0\n",
+    ")  # TODO: reshape\n",
     "with torch.no_grad():\n",
     "    if USE_TORCH_DIFFEQ:\n",
     "        traj = torchdiffeq.odeint(\n",
@@ -11978,7 +11986,7 @@
     ")\n",
     "img = ToPILImage()(grid)\n",
     "plt.imshow(img)\n",
-    "cond_values = \", \".join([ f\"{float(item):.2f}\" for item in generated_class_list[0:10,0] ])\n",
+    "cond_values = \", \".join([f\"{float(item):.2f}\" for item in generated_class_list[0:10, 0]])\n",
     "plt.title(f\"float conditional cfm\\nlabels: {cond_values}\")\n",
     "plt.savefig(\"floatconditional-cfm_noninteger.svg\")"
    ]
@@ -26069,7 +26077,9 @@
     "    for i, data in enumerate(train_loader):\n",
     "        optimizer.zero_grad()\n",
     "        x1 = data[0].to(device)\n",
-    "        y = data[1].float().to(device).reshape((batch_size, 1)) / 2. #just to have a floating point label\n",
+    "        y = (\n",
+    "            data[1].float().to(device).reshape((batch_size, 1)) / 2.0\n",
+    "        )  # just to have a floating point label\n",
     "        x0 = torch.randn_like(x1)\n",
     "        t, xt, ut, _, y1 = FM.guided_sample_location_and_conditional_flow(x0, x1, y1=y)\n",
     "        vt = model(t, xt, y1)\n",
@@ -82572,7 +82582,9 @@
     "    for i, data in tqdm(enumerate(train_loader)):\n",
     "        optimizer.zero_grad()\n",
     "        x1 = data[0].to(device)\n",
-    "        y = data[1].float().to(device).reshape((batch_size, 1)) / 2. #just to have a floating point label\n",
+    "        y = (\n",
+    "            data[1].float().to(device).reshape((batch_size, 1)) / 2.0\n",
+    "        )  # just to have a floating point label\n",
     "        x0 = torch.randn_like(x1)\n",
     "        t, xt, ut, _, y1, eps = FM.guided_sample_location_and_conditional_flow(\n",
     "            x0, x1, y1=y, return_noise=True\n",

tests/test_models.py

@@ -17,34 +17,31 @@ def test_initialize_unet():
 
     batch = torch.zeros((8, 1, 28, 28), dtype=torch.float32)
     label = torch.ones((8,), dtype=torch.long)
-    timesteps = torch.linspace(0,1,steps=8)
+    timesteps = torch.linspace(0, 1, steps=8)
 
-    _ = model(t= timesteps, x=batch, y=label)
+    _ = model(t=timesteps, x=batch, y=label)
 
 
 def test_initialize_mlp():
-
     model1 = MLP(dim=2, time_varying=True, w=64)
     batch = torch.ones((8, 3), dtype=torch.float32)
     output1 = model1(x=batch)
 
-    assert output1.shape == (8,2)
+    assert output1.shape == (8, 2)
 
     model2 = MLP(dim=2, w=64)
     batch = torch.ones((8, 2), dtype=torch.float32)
     output2 = model2(x=batch)
 
-    assert output2.shape == (8,2)
+    assert output2.shape == (8, 2)
 
 
 class mock_embedding(torch.nn.Module):
-
     def __init__(self, outdim=128):
         super().__init__()
         self.outdim = outdim
 
     def forward(self, inputs):
-
         batchsize = inputs.size(0)
         if len(inputs.shape) == 1:
             inputs = inputs.reshape((batchsize, 1))
@@ -53,21 +50,20 @@ def forward(self, inputs):
 
 
 def test_conditional_model_without_integer_labels():
-
     model_channels = 32
     model = UNetModel(
         dim=(1, 28, 28),
         num_channels=model_channels,
         num_res_blocks=1,
         class_cond=False,
-        embedding_net = mock_embedding
+        embedding_net=mock_embedding,
     )
 
     x1 = torch.ones((8, 1, 28, 28), dtype=torch.float32)
     x0 = torch.randn_like(x1)
-    FM = ConditionalFlowMatcher(sigma=0.)
+    FM = ConditionalFlowMatcher(sigma=0.0)
     t, xt, ut = FM.sample_location_and_conditional_flow(x0, x1)
 
-    label = 42.1*torch.ones((8,)).float()
+    label = 42.1 * torch.ones((8,)).float()
 
     vt = model(t=t, x=xt, y=label)

torchcfm/models/unet/unet.py

@@ -8,8 +8,15 @@
 import torch.nn.functional as F
 
 from .fp16_util import convert_module_to_f16, convert_module_to_f32
-from .nn import (avg_pool_nd, checkpoint, conv_nd, linear, normalization,
-                 timestep_embedding, zero_module)
+from .nn import (
+    avg_pool_nd,
+    checkpoint,
+    conv_nd,
+    linear,
+    normalization,
+    timestep_embedding,
+    zero_module,
+)
 
 
 class AttentionPool2d(nn.Module):
@@ -411,7 +418,7 @@ def __init__(
         use_scale_shift_norm=False,
         resblock_updown=False,
         use_new_attention_order=False,
-        embedding_net = nn.Identity
+        embedding_net=nn.Identity,
     ):
         super().__init__()
 
@@ -445,8 +452,9 @@ def __init__(
         if self.num_classes is not None:
             embedding_net = nn.Embedding if embedding_net == nn.Identity else embedding_net
             self.label_emb = embedding_net(num_classes, time_embed_dim)
-            assert not isinstance(self.label_emb, nn.Identity), f"for class-conditional networks, provide an embedding please!"
-
+            assert not isinstance(
+                self.label_emb, nn.Identity
+            ), f"for class-conditional networks, provide an embedding please!"
 
         ch = input_ch = int(channel_mult[0] * model_channels)
         self.input_blocks = nn.ModuleList(
@@ -605,9 +613,9 @@ def forward(self, t, x, y=None):
         :return: an [N x C x ...] Tensor of outputs.
         """
         timesteps = t
-        #assert (y is not None) == (
+        # assert (y is not None) == (
         #    self.num_classes is not None
-        #), "must specify y if and only if the model is class-conditional"
+        # ), "must specify y if and only if the model is class-conditional"
         while timesteps.dim() > 1:
             print(timesteps.shape)
             timesteps = timesteps[:, 0]
@@ -623,8 +631,10 @@ def forward(self, t, x, y=None):
         # (batchsize, 4*self.model_channels)
         emb = self.time_embed(embedded_timesteps)
 
-        if (y is not None):
-            assert y.shape[0] == x.shape[0], f"batch dimension of y ({y.shape[0]}) does not match x ({x.shape[0]})"
+        if y is not None:
+            assert (
+                y.shape[0] == x.shape[0]
+            ), f"batch dimension of y ({y.shape[0]}) does not match x ({x.shape[0]})"
             labels = self.label_emb(y)
             emb = emb + labels
 
@@ -862,6 +872,7 @@ def forward(self, x, timesteps):
 
 NUM_CLASSES = 1000
 
+
 # this overwrites UNetModel in __init__.py
 class UNetModelWrapper(UNetModel):
     def __init__(
@@ -883,7 +894,7 @@ def __init__(
         resblock_updown=False,
         use_fp16=False,
         use_new_attention_order=False,
-            embedding_net = nn.Identity
+        embedding_net=nn.Identity,
     ):
         """Dim (tuple): (C, H, W)"""
         image_size = dim[-1]
@@ -927,7 +938,7 @@ def __init__(
             use_scale_shift_norm=use_scale_shift_norm,
             resblock_updown=resblock_updown,
             use_new_attention_order=use_new_attention_order,
-            embedding_net = embedding_net
+            embedding_net=embedding_net,
         )
 
     def forward(self, t, x, y=None, *args, **kwargs):