Merge pull request #10 from TalkToTheGAN/misc/visualization

Visualizer enabler code.

README.md

@@ -12,8 +12,23 @@ The idea is from paper [SeqGAN: Sequence Generative Adversarial Nets with Policy
 
 The code is rewrited in PyTorch with the structure largely from [Tensorflow Implementation](https://github.com/LantaoYu/SeqGAN)
 
-## Runing
+## Running
 ```
 $ python main.py
 ```
 After runing this file, the results will be printed on terminal. You can change the parameters in the ```main.py```.
+
+
+__Using CUDA__
+
+Pass in the gpu device number for e.g. `0`
+```
+$ python main.py --cude {GPU_DEVICE_NUMBER}
+```
+
+__Enable Visualization__
+
+Run with `--visualize` parameter
+```
+$ python main.py --cude {GPU_DEVICE_NUMBER} --visualize
+```

eval/helper.py

@@ -0,0 +1,100 @@
+# -*- coding: utf-8 -*-
+
+import torch
+import math
+irange = range
+
+
+def make_grid(tensor, nrow=8, padding=2,
+              normalize=False, range=None, scale_each=False, pad_value=0):
+    """Make a grid of images.
+
+    Args:
+        tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W)
+            or a list of images all of the same size.
+        nrow (int, optional): Number of images displayed in each row of the grid.
+            The Final grid size is (B / nrow, nrow). Default is 8.
+        padding (int, optional): amount of padding. Default is 2.
+        normalize (bool, optional): If True, shift the image to the range (0, 1),
+            by subtracting the minimum and dividing by the maximum pixel value.
+        range (tuple, optional): tuple (min, max) where min and max are numbers,
+            then these numbers are used to normalize the image. By default, min and max
+            are computed from the tensor.
+        scale_each (bool, optional): If True, scale each image in the batch of
+            images separately rather than the (min, max) over all images.
+        pad_value (float, optional): Value for the padded pixels.
+
+    """
+    if not (torch.is_tensor(tensor) or
+            (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
+        raise TypeError('tensor or list of tensors expected, got {}'.format(type(tensor)))
+
+    # if list of tensors, convert to a 4D mini-batch Tensor
+    if isinstance(tensor, list):
+        tensor = torch.stack(tensor, dim=0)
+
+    if tensor.dim() == 2:  # single image H x W
+        tensor = tensor.view(1, tensor.size(0), tensor.size(1))
+    if tensor.dim() == 3:  # single image
+        if tensor.size(0) == 1:  # if single-channel, convert to 3-channel
+            tensor = torch.cat((tensor, tensor, tensor), 0)
+        return tensor
+    if tensor.dim() == 4 and tensor.size(1) == 1:  # single-channel images
+        tensor = torch.cat((tensor, tensor, tensor), 1)
+
+    if normalize is True:
+        tensor = tensor.clone()  # avoid modifying tensor in-place
+        if range is not None:
+            assert isinstance(range, tuple), \
+                "range has to be a tuple (min, max) if specified. min and max are numbers"
+
+        def norm_ip(img, min, max):
+            img.clamp_(min=min, max=max)
+            img.add_(-min).div_(max - min)
+
+        def norm_range(t, range):
+            if range is not None:
+                norm_ip(t, range[0], range[1])
+            else:
+                norm_ip(t, t.min(), t.max())
+
+        if scale_each is True:
+            for t in tensor:  # loop over mini-batch dimension
+                norm_range(t, range)
+        else:
+            norm_range(tensor, range)
+
+    # make the mini-batch of images into a grid
+    nmaps = tensor.size(0)
+    xmaps = min(nrow, nmaps)
+    ymaps = int(math.ceil(float(nmaps) / xmaps))
+    height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding)
+    grid = tensor.new(3, height * ymaps + padding, width * xmaps + padding).fill_(pad_value)
+    k = 0
+    for y in irange(ymaps):
+        for x in irange(xmaps):
+            if k >= nmaps:
+                break
+            grid.narrow(1, y * height + padding, height - padding)\
+                .narrow(2, x * width + padding, width - padding)\
+                .copy_(tensor[k])
+            k = k + 1
+    return grid
+
+
+def save_image(tensor, filename, nrow=8, padding=2,
+               normalize=False, range=None, scale_each=False, pad_value=0):
+    """Save a given Tensor into an image file.
+
+    Args:
+        tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,
+            saves the tensor as a grid of images by calling ``make_grid``.
+        **kwargs: Other arguments are documented in ``make_grid``.
+    """
+    from PIL import Image
+    tensor = tensor.cpu()
+    grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
+                     normalize=normalize, range=range, scale_each=scale_each)
+    ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy()
+    im = Image.fromarray(ndarr)
+    im.save(filename)

main.py

@@ -25,6 +25,8 @@
 from loss import *
 
 
+isDebug = True
+
 # ================== Parameter Definition =================
 
 # Basic Training Parameters
@@ -40,16 +42,16 @@
 EVAL_FILE = 'eval.data'
 VOCAB_SIZE = 5000
 # pre-training
-PRE_EPOCH_GEN = 1 # default is 120
-PRE_EPOCH_DIS = 1 # default is 5
-PRE_ITER_DIS = 1 # default is 3
+PRE_EPOCH_GEN = 1 if isDebug else 120
+PRE_EPOCH_DIS = 1 if isDebug else 5
+PRE_ITER_DIS = 1 if isDebug else 3
 # adversarial training
 GD = 'RELAX' # REBAR or RELAX
 UPDATE_RATE = 0.8
 TOTAL_BATCH = 100
-G_STEPS = 1 # default is 1
-D_STEPS = 4 # default is 4
-D_EPOCHS = 2 # default is 2
+G_STEPS = 1 if isDebug else 1
+D_STEPS = 4 if isDebug else 4
+D_EPOCHS = 2 if isDebug else 2
 # Generator Parameters
 g_emb_dim = 32
 g_hidden_dim = 32
@@ -70,8 +72,8 @@
 
 def main(opt):
 
-    cuda = opt.cuda
-    print(cuda)
+    cuda = opt.cuda; visualize = opt.visualize
+    print(f"cuda = {cuda}, visualize = {opt.visualize}")
 
     # Define Networks
     generator = Generator(VOCAB_SIZE, g_emb_dim, g_hidden_dim, cuda)
@@ -114,7 +116,7 @@ def main(opt):
     dis_optimizer = optim.Adam(discriminator.parameters())
     if opt.cuda:
         dis_criterion = dis_criterion.cuda()
-    print('Pretrain Dsicriminator ...')
+    print('Pretrain Discriminator ...')
     for epoch in range(PRE_EPOCH_DIS):
         generate_samples(generator, BATCH_SIZE, GENERATED_NUM, NEGATIVE_FILE)
         dis_data_iter = DisDataIter(POSITIVE_FILE, NEGATIVE_FILE, BATCH_SIZE)
@@ -253,9 +255,23 @@ def main(opt):
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser(description='Training Parameter')
+    parser.add_argument('--visualize', action='store_true', help='Enables Visdom')
     parser.add_argument('--cuda', action='store', default=None, type=int)
     opt = parser.parse_args()
     if opt.cuda is not None and opt.cuda >= 0:
         torch.cuda.set_device(opt.cuda)
         opt.cuda = True if torch.cuda.is_available() else False
+
+    try:
+        from eval.helper import *
+        from eval.BLEU_score import *
+        import torchnet as tnt
+        from torchnet.engine import Engine
+        from torchnet.logger import VisdomPlotLogger, VisdomLogger
+        canVisualize = True
+    except ImportError as ie:
+        eprint("Could not import vizualization imports. ")
+        canVisualize = False
+
+    opt.visualize = True if (opt.visualize and canVisualize) else False
     main(opt)

utils.py

@@ -2,7 +2,7 @@
 Auxiliary functions used during training.
 '''
 
-import os
+import os, sys
 import random
 import math
 
@@ -25,6 +25,8 @@
 
 from main import g_sequence_len, BATCH_SIZE, VOCAB_SIZE
 
+def eprint(*args, **kwargs):
+    print(*args, file=sys.stderr, **kwargs)
 
 def generate_samples(model, batch_size, generated_num, output_file):
     samples = []