Add plot_mix_temperature.py; update README.md

Author: rgflowopen

README.md

@@ -1,3 +1,124 @@
-# RG-Flow: A hierarchical and explainable flow models based on renormalization group and sparse prior
+# RG-Flow
 
-The code requires Python >= 3.7 and PyTorch >= 1.6, with optional CUDA support.
+This repository contains the code for the paper "RG-Flow: A hierarchical and explainable flow model based on renormalization group and sparse prior".
+
+TODO: add the link and the bibtex entry to the paper
+
+# Dependencies
+
+The code requires `Python >= 3.7` and `PyTorch >= 1.6`, with optional CUDA support. Other dependencies can be installed via
+
+`pip install -r requirements.txt`
+
+# Running experiments
+
+`main.py` is the code for training the network. All adjustable arguments are stored in `args.py`, together with their default values when we were training on the CelebA datset. They can be displayed via `python main.py --help`:
+
+```
+usage: main.py [-h] [--data {celeba32,celeba64,mnist32,cifar10,chair600}] [--data_path DATA_PATH]
+               [--nchannels NCHANNELS] [--L L] [--prior {gaussian,laplace}] [--subnet {rnvp,ar}]
+               [--kernel_size KERNEL_SIZE] [--nlayers NLAYERS] [--nresblocks NRESBLOCKS]
+               [--nmlp NMLP] [--nhidden NHIDDEN] [--dtype {float32,float64}]
+               [--batch_size BATCH_SIZE] [--lr LR] [--weight_decay WEIGHT_DECAY] [--epoch EPOCH]
+               [--clip_grad CLIP_GRAD] [--no_stdout] [--print_step PRINT_STEP]
+               [--save_epoch SAVE_EPOCH] [--keep_epoch KEEP_EPOCH] [--plot_epoch PLOT_EPOCH]
+               [--cuda CUDA] [--out_infix OUT_INFIX] [-o OUT_DIR]
+
+optional arguments:
+  -h, --help            show this help message and exit
+
+dataset parameters:
+  --data {celeba32,celeba64,mnist32,cifar10,chair600}
+                        dataset name
+  --data_path DATA_PATH
+                        dataset path
+  --nchannels NCHANNELS
+                        number of channels
+  --L L                 edge length of images
+
+network parameters:
+  --prior {gaussian,laplace}
+                        prior of latent variables
+  --subnet {rnvp,ar}    type of subnet in an RG block
+  --kernel_size KERNEL_SIZE
+                        edge length of an RG block
+  --nlayers NLAYERS     number of subnet layers in an RG block
+  --nresblocks NRESBLOCKS
+                        number of residual blocks in a subnet layer
+  --nmlp NMLP           number of MLP hidden layers in an residual block
+  --nhidden NHIDDEN     width of MLP hidden layers
+  --dtype {float32,float64}
+                        dtype
+
+optimizer parameters:
+  --batch_size BATCH_SIZE
+                        batch size
+  --lr LR               learning rate
+  --weight_decay WEIGHT_DECAY
+                        weight decay
+  --epoch EPOCH         number of epoches
+  --clip_grad CLIP_GRAD
+                        global norm to clip gradients, 0 for disabled
+
+system parameters:
+  --no_stdout           do not print log to stdout, for better performance
+  --print_step PRINT_STEP
+                        number of batches to print log, 0 for disabled
+  --save_epoch SAVE_EPOCH
+                        number of epochs to save network weights, 0 for disabled
+  --keep_epoch KEEP_EPOCH
+                        number of epochs to keep saved network weights, 0 for disabled
+  --plot_epoch PLOT_EPOCH
+                        number of epochs to plot samples, 0 for disabled
+  --cuda CUDA           IDs of GPUs to use, empty for disabled
+  --out_infix OUT_INFIX
+                        infix in output filename to distinguish repeated runs
+  -o OUT_DIR, --out_dir OUT_DIR
+                        directory for output, empty for disabled
+```
+
+During training, the log file and the network weights will be saved in `out_dir`.
+
+After the network is trained, `plot_mix_temperature.py` can be used to plot samples using mixed effective temperatures, described in Appendix B of the paper.
+
+# Gallery
+
+## RG-Flow structure
+
+![RG-Flow structure](docs/structure.png)
+
+## Random walk in high-level latent representations
+
+![Random walk in high-level latent representations](docs/high_level_walk.gif)
+
+## Random walk in mid-level latent representations
+
+![Random walk in mid-level latent representations](docs/mid_level_walk.gif)
+
+## Learned receptive fields
+
+![Learned receptive fields](docs/RF.png)
+
+## Learned factors
+
+![Learned factors](docs/factors.png)
+
+### High-level factor: emotion
+
+![High-level factor: emotion](docs/smile_video.gif)
+
+### High-level factor: gender
+
+![High-level factor: gender](docs/gender_video.gif)
+
+### Mid-level factor: light direction
+
+![Mid-level factor: light direction](docs/projection_video.gif)
+
+### Mid-level factor: rotation
+
+![Mid-level factor: rotation](docs/rotation_video.gif)
+
+## Face mixing in the scaling direction
+
+![Face mixing in the scaling direction](docs/mix.png)

args.py

@@ -13,11 +13,10 @@
     default='celeba32',
     choices=['celeba32', 'celeba64', 'mnist32', 'cifar10', 'chair600'],
     help='dataset name')
-group.add_argument(
-    '--data_path',
-    type=str,
-    default='./data',    # You should put your dataset directory here
-    help='dataset path')
+group.add_argument('--data_path',
+                   type=str,
+                   default='./data',
+                   help='dataset path')
 group.add_argument('--nchannels',
                    type=int,
                    default=3,
@@ -109,7 +108,7 @@
                    '--out_dir',
                    type=str,
                    default='./saved_model',
-                   help='directory prefix for output, empty for disabled')
+                   help='directory for output, empty for disabled')
 
 args = parser.parse_args()
 

docs/RF.png

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+
+from math import sqrt
+
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+from torch.distributions.laplace import Laplace
+
+import utils
+from args import args
+from main import build_mera
+
+T_low = 0.2    # Effective temperature for low-level latent variables
+T_high = 0.8    # Effective temperature for high-level latent variables
+level_cutoff = 1    # Cutoff level (\lambda in the paper)
+
+
+def main():
+    flow = build_mera()
+    last_epoch = utils.get_last_checkpoint_step()
+    utils.load_checkpoint(last_epoch, flow)
+    flow.train(False)
+
+    shape = (16, args.nchannels, args.L, args.L)
+    prior_low = Laplace(torch.tensor(0.), torch.tensor(T_low / sqrt(2)))
+    z = prior_low.sample(shape)
+    prior_high = Laplace(torch.tensor(0.), torch.tensor(T_high / sqrt(2)))
+    z_high = prior_high.sample(shape)
+    k = 2**level_cutoff
+    z[:, :, ::k, ::k] = z_high[:, :, ::k, ::k]
+    z = z.to(args.device)
+
+    with torch.no_grad():
+        x, _ = flow.inverse(z)
+
+    samples = x.permute(0, 2, 3, 1).detach().cpu().numpy()
+    samples = 1 / (1 + np.exp(-samples))
+
+    fig, axes = plt.subplots(4, 4, figsize=(4, 4), sharex=True, sharey=True)
+    for i in range(4):
+        for j in range(4):
+            ax = axes[i, j]
+            ax.imshow(samples[j * 4 + i])
+            ax.axis('off')
+    plt.tight_layout()
+    plt.savefig('./mix_T.pdf', bbox_inches='tight')
+
+
+if __name__ == '__main__':
+    main()

docs/factors.png

@@ -1,5 +1,5 @@
-matplotlib
-numpy
-scikit-image
-torch
-torchvision
+matplotlib>=3.3
+numpy>=1.19
+scikit-image>=0.17
+torch>=1.6
+torchvision>=0.7