add aunet

Author: JavaZeroo

check_model.ipynb

@@ -294,6 +294,65 @@
     "\n",
     "test(20.3423, 3.1427)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "24M\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "from utils.bridge.models import UNetModel\n",
+    "\n",
+    "kwargs = {'in_channels': 2, \n",
+    "            'model_channels': 64, \n",
+    "            'out_channels': 1, \n",
+    "            'num_res_blocks': 4, \n",
+    "            'attention_resolutions': (0,), \n",
+    "            'dropout': 0.0, \n",
+    "            'channel_mult': (1, 2, 4), \n",
+    "            'num_classes': None, \n",
+    "            'use_checkpoint': False, \n",
+    "            'num_heads': 8, \n",
+    "            'num_heads_upsample': -1, \n",
+    "            'use_scale_shift_norm': True\n",
+    "            }\n",
+    "\n",
+    "model = UNetModel(**kwargs)\n",
+    "before_train = None\n",
+    "after_train = None\n",
+    "print(f\"{int(sum(p.numel() for p in model.parameters())/1e6)}M\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([100, 1, 28, 28])"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "x = torch.rand(100, 2, 28, 28)\n",
+    "t = torch.rand(100)\n",
+    "model(x, t).shape"
+   ]
   }
  ],
  "metadata": {

test.py

@@ -1,31 +1,24 @@
 import numpy as np
-import matplotlib.pyplot as plt
 import pickle
 
 import torch
-from torch import nn, optim
-from torch.utils.data import Dataset, DataLoader
-import torch.nn.functional as F
 
-from torch.optim.lr_scheduler import CosineAnnealingLR, CosineAnnealingWarmRestarts, StepLR, OneCycleLR
 from pathlib import Path
-from sklearn.datasets import *
 
 from rich.panel import Panel
 from rich.pretty import Pretty
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
 import time as tt
 
-from utils.Datasets import BBdataset, MNISTdataset
 from utils.utils import plot_source_and_target_mnist, binary, save_gif_frame_mnist
-from utils.data_utils import gen_mnist_data, reverse_normalize_dataset, normalize_dataset_with_metadata, gen_ds
+from utils.data_utils import gen_mnist_data, reverse_normalize_dataset, normalize_dataset_with_metadata
 from utils.model_utils import get_model_before_after
 import argparse
 
 def check_model_task(args):
     if args.task == 'gaussian2mnist':
-        assert args.model in ['tunet++', 'unet++', 'unet']
+        assert args.model in ['tunet++', 'aunet', 'unet++', 'unet']
         args.time_expand = False
     else:
         assert args.model in ['mlp', 'unet++', 'unet']
@@ -45,7 +38,7 @@ def main():
     parser.add_argument('--lr', type=float, default=1e-4)
     parser.add_argument('--iter_nums', type=int, default=1)
     parser.add_argument('--epoch_nums', type=int, default=2)
-    parser.add_argument('--batch_size', type=int, default=8000)
+    parser.add_argument('-b', '--batch_size', type=int, default=8000)
     parser.add_argument('-n','--normalize', action='store_true')
     parser.add_argument('--tarined_data', action='store_true')
     parser.add_argument('--filter_number', type=int)
@@ -60,7 +53,7 @@ def main():
     np.random.seed(seed)
 
     experiment_name = args.task 
-    if args.change_epsilons:
+    if args.change_epsilons: 
         experiment_name += '_change_epsilons'
     if args.filter_number is not None and 'mnist' in args.task:
         experiment_name += f'_filter{args.filter_number}'
@@ -138,10 +131,12 @@ def main_worker(args):
             for i in range(len(test_ts) - 1):
                 dt = (test_ts[i+1] - test_ts[i])
                 test_source_reshaped = test_source
+                
                 if args.time_expand:
                     test_ts_reshaped = test_ts[i].repeat(test_source.shape[0]).reshape(-1, 1, 1, 1).repeat(1, 1, 28, 28)
                 else:
                     test_ts_reshaped = torch.unsqueeze(test_ts[i], dim=0).T
+
                 pred_bridge_reshaped = pred_bridge[i]
 
                 ret = normalize_dataset_with_metadata(real_metadata, source=test_source_reshaped, ts=test_ts_reshaped, bridge=pred_bridge_reshaped)
@@ -156,7 +151,6 @@ def main_worker(args):
                     time = test_ts_reshaped.to(args.device)
                 if before_train is not None:
                     x = before_train(x)
-
                 x = x.to(args.device)
                 model = model.to(args.device)
                 dydt = model(x, time) if time is not None else model(x)

train.py

@@ -24,7 +24,7 @@
 
 def check_model_task(args):
     if args.task.startswith('gaussian2mnist'):
-        assert args.model in ['tunet++', 'unet++', 'unet']
+        assert args.model in ['tunet++', 'aunet']
         args.time_expand = False
     else:
         assert args.model in ['mlp', 'unet++', 'unet']
@@ -43,7 +43,7 @@ def main():
     parser.add_argument('--lr', type=float, default=1e-4)
     parser.add_argument('--iter_nums', type=int, default=1)
     parser.add_argument('--epoch_nums', type=int, default=3)
-    parser.add_argument('--batch_size', type=int, default=8000)
+    parser.add_argument('-b','--batch_size', type=int, default=8000)
     parser.add_argument('-n','--normalize', action='store_true')
     parser.add_argument('--num_workers', type=int, default=20)
     parser.add_argument('--filter_number', type=int)
@@ -80,11 +80,10 @@ def main():
 def train(args, model, train_dl, optimizer, scheduler, loss_fn, before_train=None, after_train=None):
     losses = 0
     for data in train_dl:
-        if args.model == 'tunet++':
+        if isinstance(data, list):
             training_data, time = data
         else:
-            training_data = data
-            time = None
+            training_data, time = data, None
 
         training_data = training_data.squeeze().float().cpu()
         x, y = training_data[:, :-args.dim], training_data[:, -args.dim:]
@@ -168,15 +167,19 @@ def main_worker(args):
                 progress.remove_task(task2)
                 torch.save(model.state_dict(), args.log_dir / f'model_{model.__class__.__name__}_{int(iter)}.pth')
                 progress.update(task1, advance=1, description="[red]Training whole dataset (lr: %2.5f) (loss=%2.5f)" % (cur_lr, now_loss))
-                progress.log(f"[green]sub dataset {int(iter%ds_info['nums_sub_ds'])} finished; Loss: {now_loss}")
+                progress.log("[green]sub dataset %d finished; Loss: %2.5f" % (int(iter%ds_info['nums_sub_ds']), now_loss))
+    
+    console.rule("[bold bright_green blink]Finished Training")
+    console.log("Final loss: %2.5f" % (loss_list[-1]))
     # Draw loss curve
     fig, ax = plt.subplots(figsize=(10, 5))
     ax.plot(loss_list)
     ax.set_title("Loss")
     fig.savefig(args.log_dir / 'loss.png')
+    console.log("Loss curve saved to {}".format(args.log_dir / 'loss.png'))
 
     torch.save(model.state_dict(), args.log_dir / f'model_{model.__class__.__name__}_final.pth')
-
+    console.log("Model saved to {}".format(args.log_dir / f'model_{model.__class__.__name__}_final.pth'))
 
 if __name__ == '__main__':
     main()

utils/bridge/__init__.py

@@ -0,0 +1,127 @@
+import copy
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+import os
+import numpy as np
+
+
+def grad_gauss(x, m, var):
+    xout = (x - m) / var
+    return -xout
+
+
+def ornstein_ulhenbeck(x, gradx, gamma):
+    xout = x + gamma * gradx + \
+        torch.sqrt(2 * gamma) * torch.randn(x.shape, device=x.device)
+    return xout
+
+
+class Langevin(torch.nn.Module):
+
+    def __init__(self, num_steps, shape, gammas, time_sampler, device=None,
+                 mean_final=torch.tensor([0., 0.]), var_final=torch.tensor([.5, .5]), mean_match=True):
+        super().__init__()
+
+        self.mean_match = mean_match
+        self.mean_final = mean_final
+        self.var_final = var_final
+
+        self.num_steps = num_steps  # num diffusion steps
+        self.d = shape  # shape of object to diffuse
+        self.gammas = gammas.float()  # schedule
+        gammas_vec = torch.ones(self.num_steps, *self.d, device=device)
+        for k in range(num_steps):
+            gammas_vec[k] = gammas[k].float()
+        self.gammas_vec = gammas_vec
+
+        if device is not None:
+            self.device = device
+        else:
+            self.device = gammas.device
+
+        self.steps = torch.arange(self.num_steps).to(self.device)
+        self.time = torch.cumsum(self.gammas, 0).to(self.device).float()
+        self.time_sampler = time_sampler
+
+    def record_init_langevin(self, init_samples):
+        mean_final = self.mean_final
+        var_final = self.var_final
+
+        x = init_samples
+        N = x.shape[0]
+        steps = self.steps.reshape((1, self.num_steps, 1)).repeat((N, 1, 1))
+        time = self.time.reshape((1, self.num_steps, 1)).repeat((N, 1, 1))
+        gammas = self.gammas.reshape((1, self.num_steps, 1)).repeat((N, 1, 1))
+        steps = time
+
+        x_tot = torch.Tensor(N, self.num_steps, *self.d).to(x.device)
+        out = torch.Tensor(N, self.num_steps, *self.d).to(x.device)
+        store_steps = self.steps
+        num_iter = self.num_steps
+        steps_expanded = time
+
+        for k in range(num_iter):
+            gamma = self.gammas[k]
+            gradx = grad_gauss(x, mean_final, var_final)
+            t_old = x + gamma * gradx
+            z = torch.randn(x.shape, device=x.device)
+            x = t_old + torch.sqrt(2 * gamma)*z
+            gradx = grad_gauss(x, mean_final, var_final)
+            t_new = x + gamma * gradx
+
+            x_tot[:, k, :] = x
+            out[:, k, :] = (t_old - t_new)  # / (2 * gamma)
+
+        return x_tot, out, steps_expanded
+
+    def record_langevin_seq(self, net, init_samples, t_batch=None, ipf_it=0, sample=False):
+        mean_final = self.mean_final
+        var_final = self.var_final
+
+        x = init_samples
+        N = x.shape[0]
+        steps = self.steps.reshape((1, self.num_steps, 1)).repeat((N, 1, 1))
+        time = self.time.reshape((1, self.num_steps, 1)).repeat((N, 1, 1))
+        gammas = self.gammas.reshape((1, self.num_steps, 1)).repeat((N, 1, 1))
+        steps = time
+
+        x_tot = torch.Tensor(N, self.num_steps, *self.d).to(x.device)
+        out = torch.Tensor(N, self.num_steps, *self.d).to(x.device)
+        store_steps = self.steps
+        steps_expanded = steps
+        num_iter = self.num_steps
+
+        if self.mean_match:
+            for k in range(num_iter):
+                gamma = self.gammas[k]
+                t_old = net(x, steps[:, k, :])
+
+                if sample & (k == num_iter-1):
+                    x = t_old
+                else:
+                    z = torch.randn(x.shape, device=x.device)
+                    x = t_old + torch.sqrt(2 * gamma) * z
+
+                t_new = net(x, steps[:, k, :])
+                x_tot[:, k, :] = x
+                out[:, k, :] = (t_old - t_new)
+        else:
+            for k in range(num_iter):
+                gamma = self.gammas[k]
+                t_old = x + net(x, steps[:, k, :])
+
+                if sample & (k == num_iter-1):
+                    x = t_old
+                else:
+                    z = torch.randn(x.shape, device=x.device)
+                    x = t_old + torch.sqrt(2 * gamma) * z
+                t_new = x + net(x, steps[:, k, :])
+
+                x_tot[:, k, :] = x
+                out[:, k, :] = (t_old - t_new)
+
+        return x_tot, out, steps_expanded
+
+    def forward(self, net, init_samples, t_batch, ipf_it):
+        return self.record_langevin_seq(net, init_samples, t_batch, ipf_it)

utils/bridge/langevin.py

@@ -0,0 +1,2 @@
+from .basic import ScoreNetwork
+from .unet import UNetModel

utils/bridge/models/__init__.py

@@ -0,0 +1 @@
+from .basic import ScoreNetwork

utils/bridge/models/basic/__init__.py

@@ -0,0 +1,37 @@
+import torch
+from .layers import MLP
+from .time_embedding import get_timestep_embedding
+
+class ScoreNetwork(torch.nn.Module):
+
+    def __init__(self, encoder_layers=[16], pos_dim=16, decoder_layers=[128,128], x_dim=2):
+        super().__init__()
+        self.temb_dim = pos_dim
+        t_enc_dim = pos_dim *2
+        self.locals = [encoder_layers, pos_dim, decoder_layers, x_dim]
+
+        self.net = MLP(2 * t_enc_dim,
+                       layer_widths=decoder_layers +[x_dim],
+                       activate_final = False,
+                       activation_fn=torch.nn.LeakyReLU())
+
+        self.t_encoder = MLP(pos_dim,
+                             layer_widths=encoder_layers +[t_enc_dim],
+                             activate_final = False,
+                             activation_fn=torch.nn.LeakyReLU())
+
+        self.x_encoder = MLP(x_dim,
+                             layer_widths=encoder_layers +[t_enc_dim],
+                             activate_final = False,
+                             activation_fn=torch.nn.LeakyReLU())
+
+    def forward(self, x, t):
+        if len(x.shape) == 1:
+            x = x.unsqueeze(0)
+
+        temb = get_timestep_embedding(t, self.temb_dim)
+        temb = self.t_encoder(temb)
+        xemb = self.x_encoder(x)
+        h = torch.cat([xemb ,temb], -1)
+        out = self.net(h) 
+        return out