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)