update

Author: phizaz

README.md

@@ -80,4 +80,57 @@ Google drive: https://drive.google.com/drive/folders/1abNP4QKGbNnymjn8607BF0cwxX
 
 ## Training
 
-Soon ...
+We provide scripts for training & evaluate DDIM and DiffAE (including latent DPM) on the following datasets: FFHQ128, FFHQ256, Bedroom128, Horse128, Celeba64 (D2C's crop).
+Usually, the evaluation results (FID's) will be available in `eval` directory.
+
+Note: Most experiment requires at least 4x V100s during training the DPM models while requiring 1x 2080Ti during training the accompanying latent DPM. 
+
+
+
+**FFHQ128**
+```
+# diffae
+python run_ffhq128.py
+# ddim
+python run_ffhq128_ddim.py
+```
+
+**FFHQ256**
+
+We only trained the DiffAE due to high computation cost.
+This requires 8x V100s.
+```
+sbatch run_ffhq256.py
+```
+
+After the task is done, you need to train the latent DPM (requiring only 1x 2080Ti)
+```
+python run_ffhq256_latent.py
+```
+
+**Bedroom128**
+
+```
+# diffae
+python run_bedroom128.py
+# ddim
+python run_bedroom128_ddim.py
+```
+
+**Horse128**
+
+```
+# diffae
+python run_horse128.py
+# ddim
+python run_horse128_ddim.py
+```
+
+**Celeba64**
+
+This experiment can be run on 2080Ti's.
+
+```
+# diffae
+python run_celeba64.py
+```

evals/ffhq128_autoenc_130M.txt

@@ -0,0 +1 @@
+{}

experiment.py

@@ -235,7 +235,6 @@ def is_last_accum(self, batch_idx):
         return (batch_idx + 1) % self.conf.accum_batches == 0
 
     def infer_whole_dataset(self,
-                            both_flips=False,
                             with_render=False,
                             T_render=None,
                             render_save_path=None):
@@ -247,25 +246,14 @@ def infer_whole_dataset(self,
             with_render: whether to also render the images corresponding to that latent
             render_save_path: lmdb output for the rendered images
         """
-        if both_flips:
-            # both original pose and its flipped version
-            data_a = self.conf.make_dataset()
-            assert not (isinstance(data_a, CelebAlmdb) and data_a.crop_d2c
-                        ), "doesn't support celeba dataset with d2c crop"
-            data_a.transform = make_transform(self.conf.img_size, flip_prob=0)
-            data_b = self.conf.make_dataset()
-            data_b.transform = make_transform(self.conf.img_size, flip_prob=1)
-            data = ConcatDataset([data_a, data_b])
+        data = self.conf.make_dataset()
+        if isinstance(data, CelebAlmdb) and data.crop_d2c:
+            # special case where we need the d2c crop
+            data.transform = make_transform(self.conf.img_size,
+                                            flip_prob=0,
+                                            crop_d2c=True)
         else:
-            data = self.conf.make_dataset()
-            if isinstance(data, CelebAlmdb) and data.crop_d2c:
-                # special case where we need the d2c crop
-                data.transform = make_transform(self.conf.img_size,
-                                                flip_prob=0,
-                                                crop_d2c=True)
-            else:
-                data.transform = make_transform(self.conf.img_size,
-                                                flip_prob=0)
+            data.transform = make_transform(self.conf.img_size, flip_prob=0)
 
         # data = SubsetDataset(data, 21)
 
@@ -690,24 +678,15 @@ def test_step(self, batch, *args, **kwargs):
 
         # it will run only one step!
         print('global step:', self.global_step)
-        # score = evaluate_lpips(sampler=self.eval_sampler,
-        #                        model=self.ema_model,
-        #                        conf=self.conf,
-        #                        device=self.device,
-        #                        val_data=self.val_data)
-        # self.log('lpips', score)
         """
         "infer" = predict the latent variables using the encoder on the whole dataset
         """
-        if 'infer' in self.conf.eval_programs or 'inferflip' in self.conf.eval_programs:
+        if 'infer' in self.conf.eval_programs:
             if 'infer' in self.conf.eval_programs:
                 print('infer ...')
-                conds = self.infer_whole_dataset(both_flips=False).float()
-                save_path = f'latent_infer/{self.conf.name}.pkl'
-            elif 'inferflip' in self.conf.eval_programs:
-                print('infer both ...')
-                conds = self.infer_whole_dataset(both_flips=True).float()
-                save_path = f'latent_infer_flip/{self.conf.name}.pkl'
+                conds = self.infer_whole_dataset().float()
+                # NOTE: always use this path for the latent.pkl files
+                save_path = f'checkpoints/{self.conf.name}/latent.pkl'
             else:
                 raise NotImplementedError()
 

main.py

@@ -0,0 +1,27 @@
+from templates import *
+from templates_latent import *
+
+if __name__ == '__main__':
+    # train the autoenc moodel
+    # this requires V100s.
+    gpus = [0, 1, 2, 3]
+    conf = bedroom128_autoenc()
+    train(conf, gpus=gpus)
+
+    # infer the latents for training the latent DPM
+    # NOTE: not gpu heavy, but more gpus can be of use!
+    gpus = [0, 1, 2, 3]
+    conf.eval_programs = ['infer']
+    train(conf, gpus=gpus, mode='eval')
+
+    # train the latent DPM
+    # NOTE: only need a single gpu
+    gpus = [0]
+    conf = bedroom128_autoenc_latent()
+    train(conf, gpus=gpus)
+
+    # unconditional sampling score
+    # NOTE: a lot of gpus can speed up this process
+    gpus = [0, 1, 2, 3]
+    conf.eval_programs = ['fid(10,10)']
+    train(conf, gpus=gpus, mode='eval')

manipulate.ipynb

@@ -0,0 +1,11 @@
+from templates import *
+from templates_latent import *
+
+if __name__ == '__main__':
+    gpus = [0, 1, 2, 3]
+    conf = bedroom128_ddpm()
+    train(conf, gpus=gpus)
+
+    gpus = [0, 1, 2, 3]
+    conf.eval_programs = ['fid10']
+    train(conf, gpus=gpus, mode='eval')

run_bedroom128.py

@@ -0,0 +1,27 @@
+from templates import *
+from templates_latent import *
+
+if __name__ == '__main__':
+    # train the autoenc moodel
+    # this can be run on 2080Ti's.
+    gpus = [0, 1, 2, 3]
+    conf = celeba64d2c_autoenc()
+    train(conf, gpus=gpus)
+
+    # infer the latents for training the latent DPM
+    # NOTE: not gpu heavy, but more gpus can be of use!
+    gpus = [0, 1, 2, 3]
+    conf.eval_programs = ['infer']
+    train(conf, gpus=gpus, mode='eval')
+
+    # train the latent DPM
+    # NOTE: only need a single gpu
+    gpus = [0]
+    conf = celeba64d2c_autoenc_latent()
+    train(conf, gpus=gpus)
+
+    # unconditional sampling score
+    # NOTE: a lot of gpus can speed up this process
+    gpus = [0, 1, 2, 3]
+    conf.eval_programs = ['fid(10,10)']
+    train(conf, gpus=gpus, mode='eval')

run_bedroom128_ddim.py

@@ -0,0 +1,27 @@
+from templates import *
+from templates_latent import *
+
+if __name__ == '__main__':
+    # train the autoenc moodel
+    # this requires V100s.
+    gpus = [0, 1, 2, 3]
+    conf = ffhq128_autoenc_130M()
+    train(conf, gpus=gpus)
+
+    # infer the latents for training the latent DPM
+    # NOTE: not gpu heavy, but more gpus can be of use!
+    gpus = [0, 1, 2, 3]
+    conf.eval_programs = ['infer']
+    train(conf, gpus=gpus, mode='eval')
+
+    # train the latent DPM
+    # NOTE: only need a single gpu
+    gpus = [0]
+    conf = ffhq128_autoenc_latent()
+    train(conf, gpus=gpus)
+
+    # unconditional sampling score
+    # NOTE: a lot of gpus can speed up this process
+    gpus = [0, 1, 2, 3]
+    conf.eval_programs = ['fid(10,10)']
+    train(conf, gpus=gpus, mode='eval')

run_celeba64.py

@@ -0,0 +1,11 @@
+from templates import *
+from templates_latent import *
+
+if __name__ == '__main__':
+    gpus = [0, 1, 2, 3]
+    conf = ffhq128_ddpm_130M()
+    train(conf, gpus=gpus)
+
+    gpus = [0, 1, 2, 3]
+    conf.eval_programs = ['fid10']
+    train(conf, gpus=gpus, mode='eval')

run_ffhq128.py

@@ -0,0 +1,10 @@
+from templates import *
+from templates_latent import *
+
+if __name__ == '__main__':
+    # 256 requires 8x v100s, in our case, on two nodes.
+    # do not run this directly, use `sbatch run_ffhq256.sh` to spawn the srun properly.
+    gpus = [0, 1, 2, 3]
+    nodes = 2
+    conf = ffhq256_autoenc()
+    train(conf, gpus=gpus, nodes=nodes)

run_ffhq128_ddim.py

@@ -0,0 +1,13 @@
+#!/bin/sh
+#SBATCH --gres=gpu:4
+#SBATCH --cpus-per-gpu=8 
+#SBATCH --mem-per-gpu=32GB
+#SBATCH --nodes=2
+#SBATCH --ntasks=8
+#SBATCH --partition=gpu-cluster
+#SBATCH --time=72:00:00
+
+export NCCL_DEBUG=INFO
+export PYTHONFAULTHANDLER=1
+
+srun python run_ffhq256.py

run_ffhq256.py

@@ -0,0 +1,18 @@
+from templates import *
+from templates_latent import *
+
+if __name__ == '__main__':
+    # do run the run_ffhq256 before using the file to train the latent DPM
+
+    # infer the latents for training the latent DPM
+    # NOTE: not gpu heavy, but more gpus can be of use!
+    gpus = [0, 1, 2, 3]
+    conf = ffhq256_autoenc()
+    conf.eval_programs = ['infer']
+    train(conf, gpus=gpus, mode='eval')
+
+    # train the latent DPM
+    # NOTE: only need a single gpu
+    gpus = [0]
+    conf = ffhq256_autoenc_latent()
+    train(conf, gpus=gpus)