StyleGAN training

scripts/gan/stylegan/README.md

@@ -2,7 +2,24 @@
 
 **Train StyleGAN**
 
-Instructions coming soon.
+First, prepare the dataset for training. Download the FFHQ dataset from (https://github.com/NVlabs/ffhq-dataset) and save it to DATASET_PATH.
+LMDB_PATH: directory to save the output dataset, N_WORKER: number of workers, DATASET_PATH: the downloaed FFHQ images folder path 
+```bash
+python prepare_data.py --out LMDB_PATH --n_worker N_WORKER DATASET_PATH
+```
+
+Second, train the StyleGAN for FFHQ dataset.
+```bash
+python train.py --path LMDB_PATH --sched
+```
+
+*Notes for training*
+1) The original tensorflow implementation can't be 100% converted by MXNet. Two functions are missing, [gradient penalty](https://github.com/NVlabs/stylegan/blob/66813a32aac5045fcde72751522a0c0ba963f6f2/training/loss.py#L50) and [blur](https://github.com/NVlabs/stylegan/blob/66813a32aac5045fcde72751522a0c0ba963f6f2/training/networks_stylegan.py#L96). The lack of gradient penalty can cause mode collapse while training, so it is neccessary to tune the learning rate based on the number of GPUs and apply early stop. The lack of blur function results in the low image quality and this is one of the important reasons that high-resolution images can't be generated via our implementation. 
+2) The training of StyleGAN is not stable at this moment due to the aforementioned reasons. We've tested the training by using 8 K80 GPUs and single GPU. Single GPU can be problematic. The following images are generated by a model trained with 8 K80 GPUs. 
+3) It takes around 4 days with 8 K80 GPUs to train a StyleGAN to generate 128x128 images.  
+
+![images](sample_train.png "Generated 128x128 FFHQ images from the trained StyleGAN")
+
 
 **Test StyleGAN**
 

scripts/gan/stylegan/demo_stylegan.py

@@ -50,7 +50,8 @@ def save_image(data, file, normalize=True, img_range=None):
         img_range = [min(data), max(data)]
 
     norm_img = normalize_image(data, img_range[0], img_range[1])
-    img = nd.clip(norm_img * 255 + 0.5, 0, 255).asnumpy().astype(np.uint8)
+    img = nd.clip(norm_img * 255 + 0.5, 0, 255).asnumpy().astype(np.uint8) 
+
     img = Image.fromarray(np.transpose(img, (1, 2, 0)))
     img.save(file)
 
@@ -61,16 +62,18 @@ def save_image(data, file, normalize=True, img_range=None):
     parser.add_argument('--n_sample', type=int, default=10, help='number of rows of sample matrix')
     parser.add_argument('--gpu_id', type=str, default='0', help='gpu id: e.g. 0. use -1 for CPU')
     parser.add_argument('--out_dir', type=str, default='samples/', help='output directory for samples')
-    parser.add_argument('--path', type=str, default='./stylegan-ffhq-1024px-new.params',
+    parser.add_argument('--path', type=str, default='./stylegan-ffhq-1024px-new.params', 
                         help='path to checkpoint file')
+
+    args = parser.parse_args()   
 
-    args = parser.parse_args()
     if args.gpu_id == '-1':
         device = mx.cpu()
     else:
         device = mx.gpu(int(args.gpu_id.strip()))
 
-    generator = StyledGenerator(code_dim=512)
+    generator = StyledGenerator(512, blur=True)
+
     generator.initialize()
     generator.collect_params().reset_ctx(device)
     generator.load_parameters(args.path, ctx=device)
@@ -82,7 +85,9 @@ def save_image(data, file, normalize=True, img_range=None):
     imgs = sample(generator, step, mean_style, args.n_sample, device)
 
     if not os.path.isdir(args.out_dir):
-        os.makedirs(args.out_dir)
+        os.makedirs(args.out_dir) 
 
     for i in range(args.n_sample):
         save_image(imgs[i], os.path.join(args.out_dir, 'sample_{}.png'.format(i)), normalize=True, img_range=(-1, 1))
+
+