Training DualGAN have some problem
Minxiangliu opened this issue · 0 comments
Minxiangliu commented
I is machine learning noobs, I want to reproduce this DualGAN example with using convolution method.
I try to rewrite part of the source code, but the training results are different from the original code.
The results can't translate images to their opposite domain.
Does anyone can help me? Many thanks!!
My code:
from __future__ import print_function, division
import scipy
import logging
logging.getLogger('tensorflow').disabled = True
from keras_contrib.layers.normalization.instancenormalization import InstanceNormalization
from keras.datasets import mnist
from keras.layers import Input, Dense, Dropout, Activation, Concatenate
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.convolutional import Conv2D, UpSampling2D
from keras.models import Sequential, Model
from keras.optimizers import Adam
from keras.initializers import glorot_uniform
import keras.backend as K
import matplotlib.pyplot as plt
import sys, os, cv2
import numpy as np
import warnings
warnings.simplefilter('ignore')
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
class DUALGAN():
def __init__(self):
self.img_rows = 32
self.img_cols = 32
self.channels = 1
self.img_shape = (self.img_rows, self.img_cols, self.channels)
# Number of filters in the first layer of G and D
self.gf = 32
self.df = 64
# Calculate output shape of D (PatchGAN)
patch = int(self.img_rows / 2**4)
self.disc_patch = (patch, patch, 1)
optimizer = Adam(0.0002, 0.5)
# Build and compile the discriminators
self.D_A = self.build_discriminator()
self.D_A.compile(loss=self.wasserstein_loss,
optimizer=optimizer,
metrics=['accuracy'])
self.D_B = self.build_discriminator()
self.D_B.compile(loss=self.wasserstein_loss,
optimizer=optimizer,
metrics=['accuracy'])
#-------------------------
# Construct Computational
# Graph of Generators
#-------------------------
# Build the generators
self.G_AB = self.build_generator()
self.G_BA = self.build_generator()
# For the combined model we will only train the generators
self.D_A.trainable = False
self.D_B.trainable = False
# The generator takes images from their respective domains as inputs
imgs_A = Input(shape=(self.img_shape))
imgs_B = Input(shape=(self.img_shape))
# Generators translates the images to the opposite domain
fake_B = self.G_AB(imgs_A)
fake_A = self.G_BA(imgs_B)
# The discriminators determines validity of translated images
valid_A = self.D_A(fake_A)
valid_B = self.D_B(fake_B)
# Generators translate the images back to their original domain
recov_A = self.G_BA(fake_B)
recov_B = self.G_AB(fake_A)
# The combined model (stacked generators and discriminators)
self.combined = Model(inputs=[imgs_A, imgs_B], outputs=[valid_A, valid_B, recov_A, recov_B])
self.combined.compile(loss=[self.wasserstein_loss, self.wasserstein_loss, 'mae', 'mae'],
optimizer=optimizer,
loss_weights=[1, 1, 100, 100])
def build_generator(self):
"""U-Net Generator"""
def conv2d(layer_input, filters, f_size=4):
"""Layers used during downsampling"""
d = Conv2D(filters, kernel_size=f_size, strides=2, padding='same')(layer_input)
d = LeakyReLU(alpha=0.2)(d)
d = InstanceNormalization()(d)
return d
def deconv2d(layer_input, skip_input, filters, f_size=4, dropout_rate=0):
"""Layers used during upsampling"""
u = UpSampling2D(size=2)(layer_input)
u = Conv2D(filters, kernel_size=f_size, strides=1, padding='same', activation='relu')(u)
if dropout_rate:
u = Dropout(dropout_rate)(u)
u = InstanceNormalization()(u)
u = Concatenate()([u, skip_input])
return u
# Image input
d0 = Input(shape=self.img_shape)
# Downsampling
d1 = conv2d(d0, self.gf)
d2 = conv2d(d1, self.gf*2)
d3 = conv2d(d2, self.gf*4)
d4 = conv2d(d3, self.gf*8)
# Upsampling
u1 = deconv2d(d4, d3, self.gf*4)
u2 = deconv2d(u1, d2, self.gf*2)
u3 = deconv2d(u2, d1, self.gf)
u4 = UpSampling2D(size=2)(u3)
output_img = Conv2D(self.channels, kernel_size=4, strides=1, padding='same', activation='tanh')(u4)
return Model(d0, output_img)
def build_discriminator(self):
def d_layer(layer_input, filters, f_size=4, normalization=True):
"""Discriminator layer"""
d = Conv2D(filters, kernel_size=f_size, strides=2, padding='same')(layer_input)
d = LeakyReLU(alpha=0.2)(d)
if normalization:
d = InstanceNormalization()(d)
return d
img = Input(shape=self.img_shape)
d1 = d_layer(img, self.df, normalization=False)
d2 = d_layer(d1, self.df*2)
d3 = d_layer(d2, self.df*4)
d4 = d_layer(d3, self.df*8)
validity = Conv2D(self.channels, kernel_size=4, strides=2, padding='same')(d4)
return Model(img, validity)
def sample_generator_input(self, X, batch_size):
# Sample random batch of images from X
idx = np.random.randint(0, X.shape[0], batch_size)
return X[idx]
def wasserstein_loss(self, y_true, y_pred):
return K.mean(np.absolute(y_pred - y_true))
def train(self, epochs, batch_size=128, sample_interval=50):
# Load the dataset
(X_train, _), (_, _) = mnist.load_data()
img_ = []
for img in X_train:
img_.append(scipy.misc.imresize(img, (32,32)))
X_train = np.array(img_)
X_train = np.expand_dims(X_train, 3)
# Rescale -1 to 1
X_train = (X_train.astype(np.float32) - 127.5) / 127.5
# Domain A and B (rotated)
X_A = X_train[:int(X_train.shape[0]/2)]
X_B = scipy.ndimage.interpolation.rotate(X_train[int(X_train.shape[0]/2):], 90, axes=(1, 2))
clip_value = 0.01
n_critic = 4
# Adversarial ground truths
valid = np.ones((batch_size,) + self.disc_patch)
fake = np.zeros((batch_size,) + self.disc_patch)
for epoch in range(epochs):
# Train the discriminator for n_critic iterations
for _ in range(n_critic):
# ----------------------
# Train Discriminators
# ----------------------
# Sample generator inputs
imgs_A = self.sample_generator_input(X_A, batch_size)
imgs_B = self.sample_generator_input(X_B, batch_size)
# Translate images to their opposite domain
fake_B = self.G_AB.predict(imgs_A)
fake_A = self.G_BA.predict(imgs_B)
# Train the discriminators
D_A_loss_real = self.D_A.train_on_batch(imgs_A, valid)
D_A_loss_fake = self.D_A.train_on_batch(fake_A, fake)
D_B_loss_real = self.D_B.train_on_batch(imgs_B, valid)
D_B_loss_fake = self.D_B.train_on_batch(fake_B, fake)
D_A_loss = 0.5 * np.add(D_A_loss_real, D_A_loss_fake)
D_B_loss = 0.5 * np.add(D_B_loss_real, D_B_loss_fake)
# Clip discriminator weights
for d in [self.D_A, self.D_B]:
for l in d.layers:
weights = l.get_weights()
weights = [np.clip(w, -clip_value, clip_value) for w in weights]
l.set_weights(weights)
# ------------------
# Train Generators
# ------------------
# Train the generators
g_loss = self.combined.train_on_batch([imgs_A, imgs_B], [valid, valid, imgs_A, imgs_B])
# Plot the progress
print ("%d [D1 loss: %f] [D2 loss: %f] [G loss: %f]" \
% (epoch, D_A_loss[0], D_B_loss[0], g_loss[0]))
# If at save interval => save generated image samples
if epoch % sample_interval == 0:
self.save_imgs(epoch, X_A, X_B)
def save_imgs(self, epoch, X_A, X_B):
r, c = 4, 4
# Sample generator inputs
imgs_A = self.sample_generator_input(X_A, c)
imgs_B = self.sample_generator_input(X_B, c)
# Images translated to their opposite domain
fake_B = self.G_AB.predict(imgs_A)
fake_A = self.G_BA.predict(imgs_B)
gen_imgs = np.concatenate([imgs_A, fake_B, imgs_B, fake_A])
gen_imgs = gen_imgs.reshape((r, c, self.img_rows, self.img_cols, 1))
# Rescale images 0 - 1
gen_imgs = 0.5 * gen_imgs + 0.5
fig, axs = plt.subplots(r, c)
cnt = 0
for i in range(r):
for j in range(c):
axs[i,j].imshow(gen_imgs[i, j, :,:,0], cmap='gray')
axs[i,j].axis('off')
cnt += 1
fig.savefig("2Dimages/mnist_%d.png" % epoch)
plt.close()
if __name__ == '__main__':
K.clear_session()
gan = DUALGAN()
gan.train(epochs=1000, batch_size=500, sample_interval=50)
Results:
