├── dataset └── .gitkeep ├── image └── .gitkeep ├── model └── .gitkeep ├── .gitignore ├── README.md ├── src ├── convert_dataset.py ├── generate.py ├── net.py └── train.py └── LICENSE /dataset/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /image/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /model/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | dataset/* 3 | !dataset/.gitkeep 4 | image/* 5 | !image/.gitkeep 6 | model/* 7 | !model/.gitkeep 8 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # chainer-vae-gan 2 | 3 | Implementation of [Autoencoding beyond pixels using a learned similarity metric](https://arxiv.org/abs/1512.09300) 4 | 5 | ## Usage 6 | 7 | ### Make dataset file 8 | 9 | ``` 10 | $ python src/convert_dataset.py image_dir images.pkl 11 | ``` 12 | 13 | ### Train model 14 | 15 | ``` 16 | $ python src/train.py -g 0 -d images.pkl -o model/test --out_image_dir image/out 17 | ``` 18 | 19 | ## License 20 | 21 | MIT License 22 | -------------------------------------------------------------------------------- /src/convert_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import cPickle as pickle 4 | 5 | in_dir = sys.argv[1] 6 | out_file = sys.argv[2] 7 | files = os.listdir(in_dir) 8 | 9 | images = [] 10 | for file_name in files[:20000]: 11 | name, ext = os.path.splitext(file_name) 12 | if not ext in ['.jpg', '.jpeg', '.png', '.gif']: 13 | continue 14 | with open(os.path.join(in_dir, file_name), 'rb') as f: 15 | images.append(f.read()) 16 | 17 | with open(out_file, 'wb') as f: 18 | pickle.dump(images, f) 19 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 Daiki Sanno 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /src/generate.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import io 4 | import os 5 | from PIL import Image 6 | import cPickle as pickle 7 | 8 | import chainer 9 | from chainer import cuda, Variable, optimizers, serializers 10 | import chainer.functions as F 11 | import chainer.links as L 12 | import net 13 | 14 | parser = argparse.ArgumentParser(description='DCGAN trainer for ETL9') 15 | parser.add_argument('--gpu', '-g', default=-1, type=int, 16 | help='GPU ID (negative value indicates CPU)') 17 | parser.add_argument('--input', '-i', required=True, type=str, 18 | help='input model file path without extension') 19 | parser.add_argument('--output', '-o', required=True, type=str, 20 | help='output image file path without extension') 21 | parser.add_argument('--dataset', '-d', default='dataset/images.pkl', type=str, 22 | help='dataset file path') 23 | parser.add_argument('--size', '-s', default=64, type=int, choices=[48, 64, 80, 96, 112, 128], 24 | help='image size') 25 | args = parser.parse_args() 26 | 27 | image_size = args.size 28 | enc = net.Encoder(density=4, size=image_size) 29 | gen = net.Generator(density=4, size=image_size) 30 | 31 | serializers.load_hdf5(args.input + '.enc.model', enc) 32 | serializers.load_hdf5(args.input + '.gen.model', gen) 33 | 34 | with open(args.dataset, 'rb') as f: 35 | images = pickle.load(f) 36 | 37 | gpu_device = None 38 | if args.gpu >= 0: 39 | cuda.check_cuda_available() 40 | gpu_device = args.gpu 41 | 42 | latent_size = gen.latent_size 43 | out_image_row_num = 10 44 | out_image_col_num = 10 45 | 46 | if gpu_device == None: 47 | enc.to_cpu() 48 | gen.to_cpu() 49 | xp = np 50 | else: 51 | enc.to_gpu(gpu_device) 52 | gen.to_gpu(gpu_device) 53 | xp = cuda.cupy 54 | 55 | x = np.zeros((out_image_row_num * 2, 3, image_size, image_size), dtype=np.float32) 56 | perm = np.random.permutation(len(images)) 57 | for j, p in enumerate(perm[:out_image_row_num]): 58 | image = images[p] 59 | offset_x = np.random.randint(8) + 13 60 | offset_y = np.random.randint(8) + 33 61 | w = 144 62 | h = 144 63 | with io.BytesIO(image) as b: 64 | pixels = np.asarray(Image.open(b).convert('RGB').crop((offset_x, offset_y, offset_x + w, offset_y + h)).resize((image_size, image_size))) 65 | pixels = pixels.astype(np.float32).transpose((2, 0, 1)).reshape((3, image_size, image_size)) 66 | x[j * 2] = pixels / 127.5 - 1 67 | x[j * 2 + 1] = pixels[:,:,::-1] / 127.5 - 1 68 | z0, mean, var = enc(Variable(xp.asarray(x), volatile=True), train=False) 69 | z = xp.zeros((out_image_row_num * out_image_col_num, latent_size)).astype(np.float32) 70 | for j in range(out_image_row_num): 71 | for k in range(1, out_image_col_num): 72 | w = np.float32((k - 1)) / (out_image_col_num - 2) 73 | z[j * out_image_col_num + k] = z0.data[j * 2] * (1 - w) + z0.data[j * 2 + 1] * w 74 | y = gen(Variable(z, volatile=True), train=False) 75 | image = cuda.to_cpu(y.data) 76 | for j in range(out_image_row_num): 77 | image[j * out_image_col_num] = x[j * 2] 78 | image = ((image + 1) * 128).clip(0, 255).astype(np.uint8) 79 | image = image.reshape((out_image_row_num, out_image_col_num, 3, image_size, image_size)).transpose((0, 3, 1, 4, 2)).reshape((out_image_row_num * image_size, out_image_col_num * image_size, 3)) 80 | Image.fromarray(image).save(args.output) 81 | -------------------------------------------------------------------------------- /src/net.py: -------------------------------------------------------------------------------- 1 | import math 2 | import numpy as np 3 | import chainer 4 | import chainer.functions as F 5 | import chainer.links as L 6 | from chainer import cuda, Variable 7 | 8 | class Encoder(chainer.Chain): 9 | def __init__(self, density=1, size=64, latent_size=100, channel=3): 10 | assert(size % 16 == 0) 11 | initial_size = size / 16 12 | super(Encoder, self).__init__( 13 | dc1 = L.Convolution2D(channel, 32 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * channel * density)), 14 | dc2 = L.Convolution2D(32 * density, 64 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 32 * density)), 15 | norm2 = L.BatchNormalization(64 * density), 16 | dc3 = L.Convolution2D(64 * density, 128 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 64 * density)), 17 | norm3 = L.BatchNormalization(128 * density), 18 | dc4 = L.Convolution2D(128 * density, 256 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 128 * density)), 19 | norm4 = L.BatchNormalization(256 * density), 20 | mean = L.Linear(initial_size * initial_size * 256 * density, latent_size, wscale=0.02 * math.sqrt(initial_size * initial_size * 256 * density)), 21 | var = L.Linear(initial_size * initial_size * 256 * density, latent_size, wscale=0.02 * math.sqrt(initial_size * initial_size * 256 * density)), 22 | ) 23 | 24 | def __call__(self, x, train=True): 25 | xp = cuda.get_array_module(x.data) 26 | h1 = F.leaky_relu(self.dc1(x)) 27 | h2 = F.leaky_relu(self.norm2(self.dc2(h1), test=not train)) 28 | h3 = F.leaky_relu(self.norm3(self.dc3(h2), test=not train)) 29 | h4 = F.leaky_relu(self.norm4(self.dc4(h3), test=not train)) 30 | mean = self.mean(h4) 31 | var = self.var(h4) 32 | rand = xp.random.normal(0, 1, var.data.shape).astype(np.float32) 33 | z = mean + F.exp(var) * Variable(rand, volatile=not train) 34 | return (z, mean, var) 35 | 36 | class Generator(chainer.Chain): 37 | def __init__(self, density=1, size=64, latent_size=100, channel=3): 38 | assert(size % 16 == 0) 39 | initial_size = size / 16 40 | super(Generator, self).__init__( 41 | g1 = L.Linear(latent_size, initial_size * initial_size * 256 * density, wscale=0.02 * math.sqrt(latent_size)), 42 | norm1 = L.BatchNormalization(initial_size * initial_size * 256 * density), 43 | g2 = L.Deconvolution2D(256 * density, 128 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 256 * density)), 44 | norm2 = L.BatchNormalization(128 * density), 45 | g3 = L.Deconvolution2D(128 * density, 64 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 128 * density)), 46 | norm3 = L.BatchNormalization(64 * density), 47 | g4 = L.Deconvolution2D(64 * density, 32 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 64 * density)), 48 | norm4 = L.BatchNormalization(32 * density), 49 | g5 = L.Deconvolution2D(32 * density, channel, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 32 * density)), 50 | ) 51 | self.density = density 52 | self.latent_size = latent_size 53 | self.initial_size = initial_size 54 | 55 | def __call__(self, z, train=True): 56 | h1 = F.reshape(F.relu(self.norm1(self.g1(z), test=not train)), (z.data.shape[0], 256 * self.density, self.initial_size, self.initial_size)) 57 | h2 = F.relu(self.norm2(self.g2(h1), test=not train)) 58 | h3 = F.relu(self.norm3(self.g3(h2), test=not train)) 59 | h4 = F.relu(self.norm4(self.g4(h3), test=not train)) 60 | return F.tanh(self.g5(h4)) 61 | 62 | class Discriminator(chainer.Chain): 63 | def __init__(self, density=1, size=64, channel=3): 64 | assert(size % 16 == 0) 65 | initial_size = size / 16 66 | super(Discriminator, self).__init__( 67 | dc1 = L.Convolution2D(channel, 32 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * channel * density)), 68 | dc2 = L.Convolution2D(32 * density, 64 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 32 * density)), 69 | norm2 = L.BatchNormalization(64 * density), 70 | dc3 = L.Convolution2D(64 * density, 128 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 64 * density)), 71 | norm3 = L.BatchNormalization(128 * density), 72 | dc4 = L.Convolution2D(128 * density, 256 * density, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 128 * density)), 73 | norm4 = L.BatchNormalization(256 * density), 74 | dc5 = L.Linear(initial_size * initial_size * 256 * density, 2, wscale=0.02 * math.sqrt(initial_size * initial_size * 256 * density)), 75 | ) 76 | 77 | def __call__(self, x, train=True): 78 | h1 = F.leaky_relu(self.dc1(x)) 79 | h2 = F.leaky_relu(self.norm2(self.dc2(h1), test=not train)) 80 | h3 = F.leaky_relu(self.norm3(self.dc3(h2), test=not train)) 81 | h4 = F.leaky_relu(self.norm4(self.dc4(h3), test=not train)) 82 | return (self.dc5(h4), h3) 83 | 84 | class Generator48(chainer.Chain): 85 | def __init__(self): 86 | latent_size = 100 87 | super(Generator48, self).__init__( 88 | g1 = L.Linear(latent_size * 2, 6 * 6 * 256, wscale=0.02 * math.sqrt(latent_size)), 89 | norm1 = L.BatchNormalization(6 * 6 * 256), 90 | g2 = L.Deconvolution2D(256, 128, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 256)), 91 | norm2 = L.BatchNormalization(128), 92 | g3 = L.Deconvolution2D(128, 64, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 128)), 93 | norm3 = L.BatchNormalization(64), 94 | g4 = L.Deconvolution2D(64, 1, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 64)), 95 | ) 96 | self.latent_size = latent_size 97 | 98 | def __call__(self, (z, y), train=True): 99 | h1 = F.reshape(F.relu(self.norm1(self.g1(z), test=not train)), (z.data.shape[0], 256, 6, 6)) 100 | h2 = F.relu(self.norm2(self.g2(h1), test=not train)) 101 | h3 = F.relu(self.norm3(self.g3(h2), test=not train)) 102 | h4 = F.sigmoid(self.g4(h3)) 103 | return h4 104 | 105 | class Discriminator48(chainer.Chain): 106 | def __init__(self): 107 | super(Discriminator48, self).__init__( 108 | dc1 = L.Convolution2D(1, 64, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4)), 109 | norm1 = L.BatchNormalization(64), 110 | dc2 = L.Convolution2D(64, 128, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 64)), 111 | norm2 = L.BatchNormalization(128), 112 | dc3 = L.Convolution2D(128, 256, 4, stride=2, pad=1, wscale=0.02 * math.sqrt(4 * 4 * 128)), 113 | norm3 = L.BatchNormalization(256), 114 | dc4 = L.Linear(6 * 6 * 256, 2, wscale=0.02 * math.sqrt(6 * 6 * 256)), 115 | ) 116 | 117 | def __call__(self, x, train=True): 118 | h1 = F.leaky_relu(self.norm1(self.dc1(x), test=not train)) 119 | h2 = F.leaky_relu(self.norm2(self.dc2(h1), test=not train)) 120 | h3 = F.leaky_relu(self.norm3(self.dc3(h2), test=not train)) 121 | return self.dc4(h3) 122 | -------------------------------------------------------------------------------- /src/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | import io 4 | import os 5 | from PIL import Image 6 | import cPickle as pickle 7 | 8 | import chainer 9 | from chainer import cuda, Variable, optimizers, serializers 10 | import chainer.functions as F 11 | import chainer.links as L 12 | import net 13 | 14 | parser = argparse.ArgumentParser(description='VAE and DCGAN trainer') 15 | parser.add_argument('--gpu', '-g', default=-1, type=int, 16 | help='GPU ID (negative value indicates CPU)') 17 | parser.add_argument('--input', '-i', default=None, type=str, 18 | help='input model file path without extension') 19 | parser.add_argument('--output', '-o', required=True, type=str, 20 | help='output model file path without extension') 21 | parser.add_argument('--iter', default=100, type=int, 22 | help='number of iteration') 23 | parser.add_argument('--out_image_dir', default=None, type=str, 24 | help='output directory to output images') 25 | parser.add_argument('--dataset', '-d', default='dataset/images.pkl', type=str, 26 | help='dataset file path') 27 | parser.add_argument('--size', '-s', default=64, type=int, choices=[48, 64, 80, 96, 112, 128], 28 | help='image size') 29 | args = parser.parse_args() 30 | 31 | image_size = args.size 32 | enc_model = net.Encoder(density=4, size=image_size) 33 | gen_model = net.Generator(density=4, size=image_size) 34 | dis_model = net.Discriminator(density=4, size=image_size) 35 | 36 | optimizer_enc = optimizers.Adam(alpha=0.0001, beta1=0.5) 37 | optimizer_enc.setup(enc_model) 38 | optimizer_enc.add_hook(chainer.optimizer.WeightDecay(0.00001)) 39 | optimizer_gen = optimizers.Adam(alpha=0.0001, beta1=0.5) 40 | optimizer_gen.setup(gen_model) 41 | optimizer_gen.add_hook(chainer.optimizer.WeightDecay(0.00001)) 42 | optimizer_dis = optimizers.Adam(alpha=0.0001, beta1=0.5) 43 | optimizer_dis.setup(dis_model) 44 | optimizer_dis.add_hook(chainer.optimizer.WeightDecay(0.00001)) 45 | 46 | if args.input != None: 47 | serializers.load_hdf5(args.input + '.enc.model', enc_model) 48 | serializers.load_hdf5(args.input + '.enc.state', optimizer_enc) 49 | serializers.load_hdf5(args.input + '.gen.model', gen_model) 50 | serializers.load_hdf5(args.input + '.gen.state', optimizer_gen) 51 | serializers.load_hdf5(args.input + '.dis.model', dis_model) 52 | serializers.load_hdf5(args.input + '.dis.state', optimizer_dis) 53 | 54 | if args.out_image_dir != None: 55 | if not os.path.exists(args.out_image_dir): 56 | try: 57 | os.mkdir(args.out_image_dir) 58 | except: 59 | print 'cannot make directory {}'.format(args.out_image_dir) 60 | exit() 61 | elif not os.path.isdir(args.out_image_dir): 62 | print 'file path {} exists but is not directory'.format(args.out_image_dir) 63 | exit() 64 | 65 | with open(args.dataset, 'rb') as f: 66 | images = pickle.load(f) 67 | 68 | gpu_device = None 69 | if args.gpu >= 0: 70 | cuda.check_cuda_available() 71 | gpu_device = args.gpu 72 | 73 | latent_size = gen_model.latent_size 74 | BATCH_SIZE = 100 75 | image_save_interval = 10000 76 | 77 | def train_one(enc, gen, dis, optimizer_enc, optimizer_gen, optimizer_dis, x_batch, gpu_device): 78 | batch_size = len(x_batch) 79 | if gpu_device == None: 80 | xp = np 81 | else: 82 | xp = cuda.cupy 83 | # encode 84 | x_in = xp.asarray(x_batch) 85 | z0, mean, var = enc(Variable(x_in)) 86 | x0 = gen(z0) 87 | y0, l0 = dis(x0) 88 | loss_enc = F.gaussian_kl_divergence(mean, var) / float(l0.data.size) 89 | loss_gen = 0 90 | loss_gen = F.softmax_cross_entropy(y0, Variable(xp.zeros(batch_size).astype(np.int32))) 91 | loss_dis = F.softmax_cross_entropy(y0, Variable(xp.ones(batch_size).astype(np.int32))) 92 | # train generator 93 | z1 = Variable(xp.random.normal(0, 1, (batch_size, latent_size)).astype(np.float32)) 94 | x1 = gen(z1) 95 | y1, l1 = dis(x1) 96 | loss_gen += F.softmax_cross_entropy(y1, Variable(xp.zeros(batch_size).astype(np.int32))) 97 | loss_dis += F.softmax_cross_entropy(y1, Variable(xp.ones(batch_size).astype(np.int32))) 98 | # train discriminator 99 | y2, l2 = dis(Variable(xp.asarray(x_batch))) 100 | loss_enc += F.mean_squared_error(l0, l2) 101 | loss_gen += 0.1 * F.mean_squared_error(l0, l2) 102 | loss_dis += F.softmax_cross_entropy(y2, Variable(xp.zeros(batch_size).astype(np.int32))) 103 | 104 | optimizer_enc.zero_grads() 105 | loss_enc.backward() 106 | optimizer_enc.update() 107 | 108 | optimizer_gen.zero_grads() 109 | loss_gen.backward() 110 | optimizer_gen.update() 111 | 112 | optimizer_dis.zero_grads() 113 | loss_dis.backward() 114 | optimizer_dis.update() 115 | 116 | return (float(loss_enc.data), float(loss_gen.data), float(loss_dis.data)) 117 | 118 | def train(enc, gen, dis, optimizer_enc, optimizer_gen, optimizer_dis, epoch_num, gpu_device=None, out_image_dir=None): 119 | if gpu_device == None: 120 | enc.to_cpu() 121 | gen.to_cpu() 122 | dis.to_cpu() 123 | xp = np 124 | else: 125 | enc.to_gpu(gpu_device) 126 | gen.to_gpu(gpu_device) 127 | dis.to_gpu(gpu_device) 128 | xp = cuda.cupy 129 | out_image_row_num = 10 130 | out_image_col_num = 10 131 | z_out_image = Variable(xp.random.uniform(-1, 1, (out_image_row_num * out_image_col_num, latent_size)).astype(np.float32)) 132 | x_batch = np.zeros((BATCH_SIZE, 3, image_size, image_size), dtype=np.float32) 133 | for epoch in xrange(1, epoch_num + 1): 134 | x_size = len(images) 135 | perm = np.random.permutation(x_size) 136 | sum_loss_enc = 0 137 | sum_loss_gen = 0 138 | sum_loss_dis = 0 139 | for i in xrange(0, x_size, BATCH_SIZE): 140 | x_batch.fill(0) 141 | for j, p in enumerate(perm[i:i + BATCH_SIZE]): 142 | image = images[p] 143 | offset_x = np.random.randint(8) + 37 144 | offset_y = np.random.randint(8) + 68 145 | w = 96 146 | h = 96 147 | with io.BytesIO(image) as b: 148 | pixels = np.asarray(Image.open(b).convert('RGB').crop((offset_x, offset_y, offset_x + w, offset_y + h)).resize((image_size, image_size))) 149 | pixels = pixels.astype(np.float32).transpose((2, 0, 1)).reshape((3, image_size, image_size)) 150 | x_batch[j] = pixels / 127.5 - 1 151 | loss_enc, loss_gen, loss_dis = train_one(enc, gen, dis, optimizer_enc, optimizer_gen, optimizer_dis, x_batch, gpu_device) 152 | sum_loss_enc += loss_enc * BATCH_SIZE 153 | sum_loss_gen += loss_gen * BATCH_SIZE 154 | sum_loss_dis += loss_dis * BATCH_SIZE 155 | if i % image_save_interval == 0: 156 | print '{} {} {}'.format(sum_loss_enc / (i + BATCH_SIZE), sum_loss_gen / (i + BATCH_SIZE), sum_loss_dis / (i + BATCH_SIZE)) 157 | if out_image_dir != None: 158 | z, m, v = enc(Variable(xp.asarray(x_batch), volatile=True), train=False) 159 | data = gen(z, train=False).data 160 | # data = gen((z_out_image), train=False).data 161 | image = ((cuda.to_cpu(data) + 1) * 128).clip(0, 255).astype(np.uint8) 162 | image = image.reshape((out_image_row_num, out_image_col_num, 3, image_size, image_size)).transpose((0, 3, 1, 4, 2)).reshape((out_image_row_num * image_size, out_image_col_num * image_size, 3)) 163 | Image.fromarray(image).save('{0}/{1:03d}_{2:07d}.png'.format(out_image_dir, epoch, i)) 164 | org_image = ((x_batch + 1) * 128).clip(0, 255).astype(np.uint8) 165 | org_image = org_image.reshape((out_image_row_num, out_image_col_num, 3, image_size, image_size)).transpose((0, 3, 1, 4, 2)).reshape((out_image_row_num * image_size, out_image_col_num * image_size, 3)) 166 | Image.fromarray(org_image).save('{0}/{1:03d}_{2:07d}_org.png'.format(out_image_dir, epoch, i)) 167 | print 'epoch: {} done'.format(epoch) 168 | print('enc loss={}'.format(sum_loss_enc / x_size)) 169 | print('gen loss={}'.format(sum_loss_gen / x_size)) 170 | print('dis loss={}'.format(sum_loss_dis / x_size)) 171 | serializers.save_hdf5('{0}_{1:03d}.enc.model'.format(args.output, epoch), enc) 172 | serializers.save_hdf5('{0}_{1:03d}.enc.state'.format(args.output, epoch), optimizer_enc) 173 | serializers.save_hdf5('{0}_{1:03d}.gen.model'.format(args.output, epoch), gen) 174 | serializers.save_hdf5('{0}_{1:03d}.gen.state'.format(args.output, epoch), optimizer_gen) 175 | serializers.save_hdf5('{0}_{1:03d}.dis.model'.format(args.output, epoch), dis) 176 | serializers.save_hdf5('{0}_{1:03d}.dis.state'.format(args.output, epoch), optimizer_dis) 177 | 178 | train(enc_model, gen_model, dis_model, optimizer_enc, optimizer_gen, optimizer_dis, args.iter, gpu_device=gpu_device, out_image_dir=args.out_image_dir) 179 | --------------------------------------------------------------------------------