def test_model_graph():
""" tensorflow.keras 中load weights不支持那个跳过不匹配的层，所以必须手动控制权重 """
yolo = keras.models.load_model(
'model_data/yolo_weights.h5') # type:keras.models.Model
tbcback = TensorBoard()
tbcback.set_model(yolo)
annotation_path = 'train.txt'
log_dir = 'logs/000/'
classes_path = 'model_data/voc_classes.txt'
anchors_path = 'model_data/yolo_anchors.txt'
class_names = get_classes(classes_path)
num_classes = len(class_names)
anchors = get_anchors(anchors_path)
input_shape = (416, 416) # multiple of 32, hw
h, w = input_shape
image_input = keras.Input(shape=(h, w, 3))
num_anchors = len(anchors)
y_true = [
keras.Input(shape=(h // {
0: 32,
1: 16,
2: 8
}[l], w // {
0: 32,
1: 16,
2: 8
}[l], num_anchors // 3, num_classes + 5)) for l in range(3)
]
model_body = yolo_body(image_input, num_anchors // 3, num_classes)
print('Create YOLOv3 model with {} anchors and {} classes.'.format(
num_anchors, num_classes))
yolo_weight = yolo.get_weights()
for i, w in enumerate(yolo_weight):
if w.shape == (1, 1, 1024, 255):
yolo_weight[i] = w[..., :(num_anchors // 3) * (num_classes + 5)]
if w.shape == (1, 1, 512, 255):
yolo_weight[i] = w[..., :(num_anchors // 3) * (num_classes + 5)]
if w.shape == (1, 1, 256, 255):
yolo_weight[i] = w[..., :(num_anchors // 3) * (num_classes + 5)]
if w.shape == (255, ):
yolo_weight[i] = w[:(num_anchors // 3) * (num_classes + 5)]
model_body.set_weights(yolo_weight)
class CartoonGAN():
def __init__(self, args):
self.model_name = 'CartoonGAN'
self.batch_size = args.batch_size
self.epochs = args.epochs
self.gpu = args.gpu_num
self.image_channels = args.image_channels
self.image_size = args.image_size
self.init_epoch = args.init_epoch
self.log_dir = args.log_dir
self.lr = args.lr
self.model_dir = args.model_dir
self.weight = args.weight
# method for generator
def generator(self):
input_shape = [self.image_size, self.image_size, self.image_channels]
input_img = Input(shape=input_shape, name="input")
# first block
x = ReflectionPadding2D(3)(input_img)
x = Conv2D(64, (7, 7),
strides=1,
use_bias=True,
padding='valid',
name="conv1")(x)
x = InstanceNormalization(name="norm1")(x)
x = Activation("relu")(x)
# down-convolution
channel = 128
for i in range(2):
x = Conv2D(channel, (3, 3),
strides=2,
use_bias=True,
padding='same',
name="conv{}_1".format(i + 2))(x)
x = Conv2D(channel, (3, 3),
strides=1,
use_bias=True,
padding='same',
name="conv{}_2".format(i + 2))(x)
x = InstanceNormalization(name="norm{}".format(i + 2))(x)
x = Activation("relu")(x)
channel = channel * 2
# residual blocks
x_res = x
for i in range(8):
x = ReflectionPadding2D(1)(x)
x = Conv2D(256, (3, 3),
strides=1,
use_bias=True,
padding='valid',
name="conv{}_1".format(i + 4))(x)
x = InstanceNormalization(name="norm{}_1".format(i + 4))(x)
x = Activation("relu")(x)
x = ReflectionPadding2D(1)(x)
x = Conv2D(256, (3, 3),
strides=1,
use_bias=True,
padding='valid',
name="conv{}_2".format(i + 4))(x)
x = InstanceNormalization(name="norm{}_2".format(i + 4))(x)
x = Add()([x, x_res])
x_res = x
# up-convolution
for i in range(2):
x = Conv2DTranspose(channel // 2,
3,
2,
padding="same",
output_padding=1,
name="deconv{}_1".format(i + 1))(x)
x = Conv2D(channel // 2, (3, 3),
strides=1,
use_bias=True,
padding="same",
name="deconv{}_2".format(i + 1))(x)
x = InstanceNormalization(name="norm_deconv" + str(i + 1))(x)
x = Activation("relu")(x)
channel = channel // 2
# last block
x = ReflectionPadding2D(3)(x)
x = Conv2D(3, (7, 7),
strides=1,
use_bias=True,
padding="valid",
name="deconv3")(x)
x = Activation("tanh")(x)
model = Model(input_img, x, name='Cartoon_Generator')
return model
# method for discriminator
def discriminator(self):
input_shape = [self.image_size, self.image_size, self.image_channels]
input_img = Input(shape=input_shape, name="input")
# first block
x = Conv2D(32, (3, 3),
strides=1,
use_bias=True,
padding='same',
name="conv1")(input_img)
x = LeakyReLU(alpha=0.2)(x)
# block loop
channel = 64
for i in range(2):
x = Conv2D(channel, (3, 3),
strides=2,
use_bias=True,
padding='same',
name="conv{}_1".format(i + 2))(x)
x = LeakyReLU(alpha=0.2)(x)
x = Conv2D(channel * 2, (3, 3),
strides=1,
use_bias=True,
padding='same',
name="conv{}_2".format(i + 2))(x)
x = InstanceNormalization()(x)
x = LeakyReLU(alpha=0.2)(x)
channel = channel * 2
# last block
x = Conv2D(256, (3, 3),
strides=1,
use_bias=True,
padding='same',
name="conv4")(x)
x = InstanceNormalization()(x)
x = LeakyReLU(alpha=0.2)(x)
x = Conv2D(1, (3, 3),
strides=1,
use_bias=True,
padding='same',
activation='sigmoid',
name="conv5")(x)
model = Model(input_img, x, name='Cartoon_Discriminator')
return model
# vgg loss function
def vgg_loss(self, y_true, y_pred):
# get vgg model
input_shape = [self.image_size, self.image_size, self.image_channels]
img_input = Input(shape=input_shape, name="vgg_input")
vgg19 = tf.keras.applications.vgg19.VGG19(weights='imagenet')
vggmodel = Model(inputs=vgg19.input,
outputs=vgg19.get_layer('block4_conv4').output)
x = vggmodel(img_input)
vgg = Model(img_input, x, name='VGG_for_Feature_Extraction')
# get l1 loss for the content loss
y_true = vgg(y_true)
y_pred = vgg(y_pred)
content_loss = tf.losses.absolute_difference(y_true, y_pred)
return content_loss
# compile each model
def compile_model(self):
# init summary writer for tensorboard
self.callback1 = TensorBoard(self.log_dir + '/discriminator')
self.callback2 = TensorBoard(self.log_dir + '/generator')
self.callback3 = TensorBoard(self.log_dir + '/generated_images')
# model stuff
input_shape = [self.image_size, self.image_size, self.image_channels]
adam1 = Adam(lr=self.lr)
adam2 = Adam(lr=self.lr * 2)
# init and add multi-gpu support
try:
self.discriminator = multi_gpu_model(self.discriminator(),
gpus=self.gpu)
except:
self.discriminator = self.discriminator()
try:
self.generator = multi_gpu_model(self.generator(), gpus=self.gpu)
except:
self.generator = self.generator()
# compile discriminator
self.discriminator.compile(loss='binary_crossentropy', optimizer=adam1)
# compile generator
input_tensor = Input(shape=input_shape)
generated_catroon_tensor = self.generator(input_tensor)
self.discriminator.trainable = False # for here we only train the generator
discriminator_output = self.discriminator(generated_catroon_tensor)
self.train_generator = Model(
input_tensor,
outputs=[generated_catroon_tensor, discriminator_output])
# add multi-gpu support
try:
self.train_generator = multi_gpu_model(self.train_generator,
gpus=self.gpu)
except:
pass
self.train_generator.compile(
loss=[self.vgg_loss, 'binary_crossentropy'],
loss_weights=[float(self.weight), 1.0],
optimizer=adam2)
# set callback model
self.callback1.set_model(self.discriminator)
self.callback2.set_model(self.train_generator)
self.callback3.set_model(self.train_generator)
# method for training process
def train(self):
# start training
flip = False
variance = 1 / 127.5
start_time = time.time()
for epoch in range(1, self.epochs + 1):
# create batch generator at each epoch
batch_generator = DataGenerator(image_size=self.image_size,
batch_size=self.batch_size)
batch_end = len(batch_generator)
print('Epoch {}'.format(epoch))
# start training for each batch
for idx, (photo, cartoon, smooth_cartoon,
index) in enumerate(batch_generator):
# these two tensors measure the output of generator and discriminator
real = np.ones((self.batch_size, ) + (64, 64, 1))
fake = np.zeros((self.batch_size, ) + (64, 64, 1))
# check if it is the end of an epoch
if index + 1 == batch_end:
break
# initial training or start training
if epoch < self.init_epoch:
g_loss = self.train_generator.train_on_batch(
photo, [photo, real])
generated_img = self.generator.predict(photo)
print(
"Batch %d (initial training for generator), g_loss: %.5f, with time: %4.4f"
% (idx, g_loss[2], time.time() - start_time))
start_time = time.time()
write_log(self.callback2, 'g_loss', g_loss[2],
idx + (epoch + 1) * len(batch_generator))
if idx % 20 == 0:
write_images(self.callback3, generated_img,
'generated_imgs',
idx + (epoch + 1) * len(batch_generator))
if epoch % 20 == 0 and K.eval(
self.train_generator.optimizer.lr) > 0.0001:
K.set_value(
self.train_generator.optimizer.lr,
K.eval(self.train_generator.optimizer.lr) * 0.99)
else:
# add noise to the input of discriminator
if variance > 0.00001:
variance = variance * 0.9999
gaussian = np.random.normal(
0, variance, (cartoon.shape[1], cartoon.shape[2]))
cartoon[:, :, :, 0] = cartoon[:, :, :, 0] + gaussian
cartoon[:, :, :, 1] = cartoon[:, :, :, 1] + gaussian
cartoon[:, :, :, 2] = cartoon[:, :, :, 2] + gaussian
gaussian = np.random.normal(
0, variance, (cartoon.shape[1], cartoon.shape[2]))
smooth_cartoon[:, :, :,
0] = smooth_cartoon[:, :, :,
0] + gaussian
smooth_cartoon[:, :, :,
1] = smooth_cartoon[:, :, :,
1] + gaussian
smooth_cartoon[:, :, :,
2] = smooth_cartoon[:, :, :,
2] + gaussian
# generate cartoonized images
generated_img = self.generator.predict(photo)
# to certain probability: flip the label of discriminator
if idx % 9 == 0 or np.random.uniform(0, 1) < 0.05:
real = fake
fake = fake + 1
flip = True
# train discriminator and adversarial loss
real_loss = self.discriminator.train_on_batch(
cartoon, real)
smooth_loss = self.discriminator.train_on_batch(
smooth_cartoon, fake)
fake_loss = self.discriminator.train_on_batch(
generated_img, fake)
d_loss = (real_loss + smooth_loss + fake_loss) / 3
# train generator
if flip:
real = fake
fake = fake - 1
flip = False
g_loss = self.train_generator.train_on_batch(
photo, [photo, real])
print(
"Batch %d, d_loss: %.5f, g_loss: %.5f, with time: %4.4f"
% (idx, d_loss, g_loss[2], time.time() - start_time))
start_time = time.time()
# add losses to writer
write_log(self.callback1, 'd_loss', d_loss,
idx + (epoch + 1) * len(batch_generator))
write_log(self.callback2, 'g_loss', g_loss[2],
idx + (epoch + 1) * len(batch_generator))
if idx % 20 == 0:
write_images(self.callback3, generated_img,
'generated_imgs',
idx + (epoch + 1) * len(batch_generator))
# change learning rate
if epoch % 20 == 0 and K.eval(
self.discriminator.optimizer.lr) > 0.0001:
K.set_value(
self.discriminator.optimizer.lr,
K.eval(self.discriminator.optimizer.lr) * 0.95)
if epoch % 20 == 0 and K.eval(
self.train_generator.optimizer.lr) > 0.0001:
K.set_value(
self.train_generator.optimizer.lr,
K.eval(self.train_generator.optimizer.lr) * 0.95)
# save model
if epoch % 50 == 0:
self.generator.save_weights(
self.model_dir + '/' +
'CartoonGan_generator_epoch_{}.h5'.format(epoch))
self.discriminator.save_weights(
self.model_dir + '/' +
'CartoonGan_discriminator_epoch_{}.h5'.format(epoch))
self.train_generator.save_weights(
self.model_dir + '/' +
'CartoonGan_train_generator_epoch_{}.h5'.format(epoch))
print('Done!')
self.generator.save('CartoonGan_generator.h5')
def train(x_train_real, input_size, batch_size, epochs, output_dir,
create_generator, create_discriminator):
pred_dir = os.path.join(output_dir, 'predictions')
model_dir = os.path.join(output_dir, 'models')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if not os.path.exists(pred_dir):
os.makedirs(pred_dir)
if not os.path.exists(model_dir):
os.makedirs(model_dir)
gan, discriminator, generator = create_gan(input_size, create_generator,
create_discriminator)
y_batch_ones_generator = np.ones([batch_size, 1])
y_batch_zeros = np.zeros([batch_size, 1])
callback = TensorBoard(log_dir=os.path.join(output_dir, 'log'))
callback.set_model(gan)
indices = np.arange(x_train_real.shape[0])
step = 0
for e in range(epochs):
print('Epoch: %d/%d' % (e + 1, epochs))
np.random.shuffle(indices)
batch_indices = [
indices[i:i + batch_size]
for i in range(0, len(indices), batch_size)
]
discriminator.save(os.path.join(model_dir, 'discriminator_%d.h5' % e))
generator.save(os.path.join(model_dir, 'generator_%d.h5' % e))
for i in tqdm(batch_indices):
# Create batch data
x_batch_real = x_train_real[i]
noise_batch = create_noise_batch(len(i), input_size)
x_batch_generator = generator.predict_on_batch(noise_batch)
# Train discriminator with label smoothing
y_batch_ones = np.random.uniform(0.8, 1, size=(len(i), 1))
discriminator_loss_m = discriminator.train_on_batch(
x_batch_real, y_batch_ones)
discriminator_loss_g = discriminator.train_on_batch(
x_batch_generator, y_batch_zeros[:len(i)])
discriminator_loss = (discriminator_loss_m +
discriminator_loss_g) / 2
# Train generator
noise_batch = create_noise_batch(len(i), input_size)
generator_loss = gan.train_on_batch(
noise_batch, y_batch_ones_generator[:len(i)])
write_log(callback,
['discriminator_train_loss', 'generator_train_loss'],
[discriminator_loss, generator_loss], step)
if step % 100 == 0:
save_predictions(generator, e, step, input_size, pred_dir)
step += 1
discriminator.save(os.path.join(model_dir, 'discriminator_%d.h5' % e))
generator.save(os.path.join(model_dir, 'generator_%d.h5' % e))
save_predictions(generator, e, step, input_size, pred_dir)
generate_gif(pred_dir, output_dir)
if __name__ == "__main__":
een = model.LatentResidualModel3Layer(opt)
base = model.BaselineModel3Layer(opt)
model_f = een.build()
base_model = base.build()
if opt.loss == 'l1':
loss = mean_absolute_error
elif opt.loss == 'l2':
loss = mean_squared_error
# load trained weight
if os.path.exists(opt.model_filename_f):
# model_g.trainable=False
print("load weight model_f ....")
model_f.load_weights(opt.model_filename_f)
optimizer = Adam(opt.lrt)
model_f.compile(optimizer=optimizer, loss=loss)
base_model.compile(optimizer=optimizer, loss=loss)
callback_f.set_model(model_f)
callback_g.set_model(base_model)
# our model
if opt.model == 'latent-3layer':
print("{} model train!....".format(opt.model_filename_f))
train(model_f, callback_f, opt.model, 500)
else:
print("{} model train!....".format(opt.model))
train(base_model, callback_g, opt.model, 500)
x = darknetconv2d_bn_leaky(x, 128, (1, 1))
x = layers.UpSampling2D(2)(x)
x = layers.Concatenate()([x, x_26])
x, y3 = make_last_layers(x, 128, num_anchors, num_classes)
h, w, _ = input_size
y1 = YoloOutputBoxLayer(anchors[6:], 1, num_classes, training)(y1)
y2 = YoloOutputBoxLayer(anchors[3:6], 2, num_classes, training)(y2)
y3 = YoloOutputBoxLayer(anchors[0:3], 3, num_classes, training)(y3)
if training:
return Model(inputs, (y1, y2, y3), name='Yolo-V3')
outputs = NMSLayer(num_classes, iou_threshold, score_threshold)([y1, y2, y3])
return Model(inputs, outputs, name='Yolo-V3')
if __name__ == "__main__":
from tensorflow.python.keras.callbacks import TensorBoard
# test yolo_v3
model = yolov3((416, 416, 3), training=True)
model.summary(line_length=250)
# test tiny yolo_v3
model = yolov3_tiny((384, 384, 3), training=True)
model.summary(line_length=250)
# Save model graph on tensorboard
model_tb = TensorBoard('../logs-2')
model_tb.set_model(model)