hello

Dol+I · Dol+I · commit a8e7d53bb7f7 · 2018-06-19T21:33:49.000+09:00
diff --git a/config.py b/config.py
@@ -0,0 +1,64 @@
+#-*- coding: utf-8 -*-
+import argparse
+
+def str2bool(v):
+    return v.lower() in ('true', '1')
+
+arg_lists = []
+parser = argparse.ArgumentParser()
+
+def add_argument_group(name):
+    arg = parser.add_argument_group(name)
+    arg_lists.append(arg)
+    return arg
+
+# Network
+net_arg = add_argument_group('Network')
+net_arg.add_argument('--input_scale_size', type=int, default=64,
+                     help='input image will be resized with the given value as width and height')
+net_arg.add_argument('--g_num_layer', type=int, default=3)
+net_arg.add_argument('--d_num_layer', type=int, default=5)
+net_arg.add_argument('--cnn_type', type=int, default=0)
+net_arg.add_argument('--fc_hidden_dim', type=int, default=128, help='only for toy dataset')
+
+# Data
+data_arg = add_argument_group('Data')
+data_arg.add_argument('--dataset', type=str, default='edges2shoes')
+data_arg.add_argument('--batch_size', type=int, default=200)
+data_arg.add_argument('--a_grayscale', type=str2bool, default=False)
+data_arg.add_argument('--b_grayscale', type=str2bool, default=False)
+data_arg.add_argument('--num_worker', type=int, default=12)
+
+# Training / test parameters
+train_arg = add_argument_group('Training')
+train_arg.add_argument('--is_train', type=str2bool, default=True)
+train_arg.add_argument('--optimizer', type=str, default='adam')
+train_arg.add_argument('--max_step', type=int, default=500000)
+train_arg.add_argument('--lr', type=float, default=0.0002)
+train_arg.add_argument('--beta1', type=float, default=0.5)
+train_arg.add_argument('--beta2', type=float, default=0.999)
+train_arg.add_argument('--loss', type=str, default="log_prob",
+                       choices=["log_prob"], help="least square loss doesn't work well")
+train_arg.add_argument('--weight_decay', type=float, default=0.0001)
+
+# Misc
+misc_arg = add_argument_group('Misc')
+misc_arg.add_argument('--load_path', type=str, default='')
+misc_arg.add_argument('--log_step', type=int, default=50)
+misc_arg.add_argument('--save_step', type=int, default=500)
+misc_arg.add_argument('--num_log_samples', type=int, default=3)
+misc_arg.add_argument('--log_level', type=str, default='INFO', choices=['INFO', 'DEBUG', 'WARN'])
+misc_arg.add_argument('--log_dir', type=str, default='../logs')
+misc_arg.add_argument('--data_dir', type=str, default='../data')
+misc_arg.add_argument('--num_gpu', type=int, default=8)
+misc_arg.add_argument('--test_data_path', type=str, default=None,
+                      help='directory with images which will be used in test sample generation')
+misc_arg.add_argument('--sample_per_image', type=int, default=64,
+                      help='# of sample per image during test sample generation')
+misc_arg.add_argument('--random_seed', type=int, default=123)
+misc_arg.add_argument('--skip_pix2pix_processing', type=str2bool, default=False,
+                      help='just for fast debugging in poor cpu machine')
+
+def get_config():
+    config, unparsed = parser.parse_known_args()
+    return config, unparsed
diff --git a/data_loader.py b/data_loader.py
@@ -0,0 +1,90 @@
+import os
+import numpy as np
+from glob import glob
+from PIL import Image
+from tqdm import tqdm
+
+import torch
+from torchvision import transforms
+
+PIX2PIX_DATASETS = [
+    'facades', 'cityscapes', 'maps', 'edges2shoes', 'edges2handbags']
+
+def makedirs(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+def pix2pix_split_images(root):
+    paths = glob(os.path.join(root, "train/*"))
+
+    a_path = os.path.join(root, "A")
+    b_path = os.path.join(root, "B")
+
+    makedirs(a_path)
+    makedirs(b_path)
+
+    for path in tqdm(paths, desc="pix2pix processing"):
+        filename = os.path.basename(path)
+
+        a_image_path = os.path.join(a_path, filename)
+        b_image_path = os.path.join(b_path, filename)
+
+        if os.path.exists(a_image_path) and os.path.exists(b_image_path):
+            continue
+
+        image = Image.open(os.path.join(path)).convert('RGB')
+        data = np.array(image)
+
+        height, width, channel = data.shape
+
+        a_image = Image.fromarray(data[:,:width/2].astype(np.uint8))
+        b_image = Image.fromarray(data[:,width/2:].astype(np.uint8))
+
+        a_image.save(a_image_path)
+        b_image.save(b_image_path)
+
+class Dataset(torch.utils.data.Dataset):
+    def __init__(self, root, scale_size, data_type, skip_pix2pix_processing=False):
+        self.root = root
+        if not os.path.exists(self.root):
+            raise Exception("[!] {} not exists.".format(root))
+
+        self.name = os.path.basename(root)
+        if self.name in PIX2PIX_DATASETS and not skip_pix2pix_processing:
+            pix2pix_split_images(self.root)
+
+        self.paths = glob(os.path.join(self.root, '{}/*'.format(data_type)))
+        if len(self.paths) == 0:
+            raise Exception("No images are found in {}".format(self.root))
+        self.shape = list(Image.open(self.paths[0]).size) + [3]
+
+        self.transform = transforms.Compose([
+            transforms.Scale(scale_size), 
+            transforms.ToTensor(), 
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ])
+
+    def __getitem__(self, index):
+        image = Image.open(self.paths[index]).convert('RGB')
+        return self.transform(image)
+
+    def __len__(self):
+        return len(self.paths)
+
+def get_loader(root, batch_size, scale_size, num_workers=2,
+               skip_pix2pix_processing=False, shuffle=True):
+    a_data_set, b_data_set = \
+        Dataset(root, scale_size, "A", skip_pix2pix_processing), \
+        Dataset(root, scale_size, "B", skip_pix2pix_processing)
+    a_data_loader = torch.utils.data.DataLoader(dataset=a_data_set,
+                                                batch_size=batch_size,
+                                                shuffle=True,
+                                                num_workers=num_workers)
+    b_data_loader = torch.utils.data.DataLoader(dataset=b_data_set,
+                                                batch_size=batch_size,
+                                                shuffle=True,
+                                                num_workers=num_workers)
+    a_data_loader.shape = a_data_set.shape
+    b_data_loader.shape = b_data_set.shape
+
+    return a_data_loader, b_data_loader
diff --git a/main.py b/main.py
@@ -0,0 +1,41 @@
+import torch
+
+from trainer import Trainer
+from config import get_config
+from data_loader import get_loader
+from utils import prepare_dirs_and_logger, save_config
+
+def main(config):
+    prepare_dirs_and_logger(config)
+
+    torch.manual_seed(config.random_seed)
+    if config.num_gpu > 0:
+        torch.cuda.manual_seed(config.random_seed)
+
+    if config.is_train:
+        data_path = config.data_path
+        batch_size = config.batch_size
+    else:
+        if config.test_data_path is None:
+            data_path = config.data_path
+        else:
+            data_path = config.test_data_path
+        batch_size = config.sample_per_image
+
+    a_data_loader, b_data_loader = get_loader(
+            data_path, batch_size, config.input_scale_size,
+            config.num_worker, config.skip_pix2pix_processing)
+
+    trainer = Trainer(config, a_data_loader, b_data_loader)
+
+    if config.is_train:
+        save_config(config)
+        trainer.train()
+    else:
+        if not config.load_path:
+            raise Exception("[!] You should specify `load_path` to load a pretrained model")
+        trainer.test()
+
+if __name__ == "__main__":
+    config, unparsed = get_config()
+    main(config)
diff --git a/models.py b/models.py
@@ -0,0 +1,113 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from torch.utils.data import TensorDataset, DataLoader
+
+class GeneratorCNN(nn.Module):
+    def __init__(self, input_channel, output_channel, conv_dims, deconv_dims, num_gpu):
+        super(GeneratorCNN, self).__init__()
+        self.num_gpu = num_gpu
+        self.layers = []
+
+        prev_dim = conv_dims[0]
+        self.layers.append(nn.Conv2d(input_channel, prev_dim, 4, 2, 1, bias=False))
+        self.layers.append(nn.LeakyReLU(0.2, inplace=True))
+
+        for out_dim in conv_dims[1:]:
+            self.layers.append(nn.Conv2d(prev_dim, out_dim, 4, 2, 1, bias=False))
+            self.layers.append(nn.BatchNorm2d(out_dim))
+            self.layers.append(nn.LeakyReLU(0.2, inplace=True))
+            prev_dim = out_dim
+
+        for out_dim in deconv_dims:
+            self.layers.append(nn.ConvTranspose2d(prev_dim, out_dim, 4, 2, 1, bias=False))
+            self.layers.append(nn.BatchNorm2d(out_dim))
+            self.layers.append(nn.ReLU(True))
+            prev_dim = out_dim
+
+        self.layers.append(nn.ConvTranspose2d(prev_dim, output_channel, 4, 2, 1, bias=False))
+        self.layers.append(nn.Tanh())
+
+        self.layer_module = nn.ModuleList(self.layers)
+
+    def main(self, x):
+        out = x
+        for layer in self.layer_module:
+            out = layer(out)
+        return out
+
+    def forward(self, x):
+        return self.main(x)
+
+class DiscriminatorCNN(nn.Module):
+    def __init__(self, input_channel, output_channel, hidden_dims, num_gpu):
+        super(DiscriminatorCNN, self).__init__()
+        self.num_gpu = num_gpu
+        self.layers = []
+
+        prev_dim = hidden_dims[0]
+        self.layers.append(nn.Conv2d(input_channel, prev_dim, 4, 2, 1, bias=False))
+        self.layers.append(nn.LeakyReLU(0.2, inplace=True))
+
+        for out_dim in hidden_dims[1:]:
+            self.layers.append(nn.Conv2d(prev_dim, out_dim, 4, 2, 1, bias=False))
+            self.layers.append(nn.BatchNorm2d(out_dim))
+            self.layers.append(nn.LeakyReLU(0.2, inplace=True))
+            prev_dim = out_dim
+
+        self.layers.append(nn.Conv2d(prev_dim, output_channel, 4, 1, 0, bias=False))
+        self.layers.append(nn.Sigmoid())
+
+        self.layer_module = nn.ModuleList(self.layers)
+
+    def main(self, x):
+        out = x
+        for layer in self.layer_module:
+            out = layer(out)
+        return out.view(out.size(0), -1)
+
+    def forward(self, x):
+        return self.main(x)
+
+class GeneratorFC(nn.Module):
+    def __init__(self, input_size, output_size, hidden_dims):
+        super(GeneratorFC, self).__init__()
+        self.layers = []
+
+        prev_dim = input_size
+        for hidden_dim in hidden_dims:
+            self.layers.append(nn.Linear(prev_dim, hidden_dim))
+            self.layers.append(nn.ReLU(True))
+            prev_dim = hidden_dim
+        self.layers.append(nn.Linear(prev_dim, output_size))
+
+        self.layer_module = nn.ModuleList(self.layers)
+
+    def forward(self, x):
+        out = x
+        for layer in self.layer_module:
+            out = layer(out)
+        return out
+
+class DiscriminatorFC(nn.Module):
+    def __init__(self, input_size, output_size, hidden_dims):
+        super(DiscriminatorFC, self).__init__()
+        self.layers = []
+
+        prev_dim = input_size
+        for idx, hidden_dim in enumerate(hidden_dims):
+            self.layers.append(nn.Linear(prev_dim, hidden_dim))
+            self.layers.append(nn.ReLU(True))
+            prev_dim = hidden_dim
+
+        self.layers.append(nn.Linear(prev_dim, output_size))
+        self.layers.append(nn.Sigmoid())
+
+        self.layer_module = nn.ModuleList(self.layers)
+
+    def forward(self, x):
+        out = x
+        for layer in self.layer_module:
+            out = layer(out)
+        return out.view(-1, 1)
diff --git a/test b/test
@@ -0,0 +1 @@
+aaa
diff --git a/trainer.py b/trainer.py
diff --git a/utils.py b/utils.py
