Ejemplo n.º 1
0
def main():
    tf.random.set_seed(233)
    np.random.seed(233)

    z_dim = 100
    epochs = 3000000
    batch_size = 512
    learning_rate = 2e-4
    # ratios = D steps:G steps
    ratios = 2

    img_path = glob.glob(os.path.join('faces', '*.jpg'))
    dataset, img_shape, _ = make_anime_dataset(img_path, batch_size)
    print(dataset, img_shape)
    sample = next(iter(dataset))
    print(sample.shape,
          tf.reduce_max(sample).numpy(),
          tf.reduce_min(sample).numpy())
    dataset = dataset.repeat()
    db_iter = iter(dataset)

    generator = Generator()
    generator.build(input_shape=(None, z_dim))
    # generator.load_weights(os.path.join('checkpoints', 'generator-5000'))
    discriminator = Discriminator()
    discriminator.build(input_shape=(None, 64, 64, 3))
    # discriminator.load_weights(os.path.join('checkpoints', 'discriminator-5000'))

    g_optimizer = tf.optimizers.Adam(learning_rate, beta_1=0.5)
    d_optimizer = tf.optimizers.Adam(learning_rate, beta_1=0.5)
    # a fixed noise for sampling
    z_sample = tf.random.normal([100, z_dim])

    g_loss_meter = keras.metrics.Mean()
    d_loss_meter = keras.metrics.Mean()
    gp_meter = keras.metrics.Mean()

    for epoch in range(epochs):

        # train D
        for step in range(ratios):
            batch_z = tf.random.normal([batch_size, z_dim])
            batch_x = next(db_iter)
            with tf.GradientTape() as tape:
                d_loss, gp = d_loss_fn(generator, discriminator, batch_z,
                                       batch_x)

            d_loss_meter.update_state(d_loss)
            gp_meter.update_state(gp)

            gradients = tape.gradient(d_loss,
                                      discriminator.trainable_variables)
            d_optimizer.apply_gradients(
                zip(gradients, discriminator.trainable_variables))

        # train G
        batch_z = tf.random.normal([batch_size, z_dim])
        with tf.GradientTape() as tape:
            g_loss = g_loss_fn(generator, discriminator, batch_z)

        g_loss_meter.update_state(g_loss)

        gradients = tape.gradient(g_loss, generator.trainable_variables)
        g_optimizer.apply_gradients(
            zip(gradients, generator.trainable_variables))

        if epoch % 100 == 0:

            fake_image = generator(z_sample, training=False)

            print(epoch, 'd-loss:',
                  d_loss_meter.result().numpy(), 'g-loss',
                  g_loss_meter.result().numpy(), 'gp',
                  gp_meter.result().numpy())

            d_loss_meter.reset_states()
            g_loss_meter.reset_states()
            gp_meter.reset_states()

            # save generated image samples
            img_path = os.path.join('images_wgan_gp', 'wgan_gp-%d.png' % epoch)
            save_result(fake_image.numpy(), 10, img_path, color_mode='P')

        if epoch + 1 % 2000 == 0:
            generator.save_weights(
                os.path.join('checkpoints_gp', 'generator-%d' % epoch))
            discriminator.save_weights(
                os.path.join('checkpoints_gp', 'discriminator-%d' % epoch))
Ejemplo n.º 2
0
def main():
    tf.random.set_seed(3333)
    np.random.seed(3333)
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    assert tf.__version__.startswith('2.')

    z_dim = 100  # 隐藏向量z的长度
    epochs = 3000000  # 训练步数
    batch_size = 64
    learning_rate = 0.0002
    is_training = True

    # 获取数据集路径
    img_path = glob.glob(r'C:\Users\jay_n\.keras\datasets\faces\*.jpg') + \
        glob.glob(r'C:\Users\jay_n\.keras\datasets\faces\*.png')
    print('images num:', len(img_path))
    # 构建数据集对象
    dataset, img_shape, _ = make_anime_dataset(img_path, batch_size, resize=64)
    print(dataset, img_shape)
    sample = next(iter(dataset))  # 采样
    print(sample.shape, tf.reduce_max(sample).numpy(), tf.reduce_min(sample).numpy())
    dataset = dataset.repeat(100)
    db_iter = iter(dataset)

    generator = Generator()
    generator.build(input_shape=(4, z_dim))
    discriminator = Discriminator()
    discriminator.build(input_shape=(4, 64, 64, 3))
    # 分别为生成器和判别器创建优化器
    g_optimizer = keras.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5)
    d_optimizer = keras.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5)

    # generator.load_weights('generator.ckpt')
    # discriminator.load_weights('discriminator.ckpt')
    # print('Loaded ckpt!!')

    d_losses, g_losses = [], []
    for epoch in range(epochs):
        # 1. 训练判别器
        for _ in range(1):
            # 采样隐藏向量
            batch_z = tf.random.normal([batch_size, z_dim])
            batch_x = next(db_iter)  # 采样真实图片
            # 判别器前向计算
            with tf.GradientTape() as tape:
                d_loss, _ = d_loss_fn(generator, discriminator, batch_z, batch_x, is_training)
            grads = tape.gradient(d_loss, discriminator.trainable_variables)
            d_optimizer.apply_gradients(zip(grads, discriminator.trainable_variables))
        # 2. 训练生成器
        # 采样隐藏向量
        batch_z = tf.random.normal([batch_size, z_dim])
        # 生成器前向计算
        with tf.GradientTape() as tape:
            g_loss = g_loss_fn(generator, discriminator, batch_z, is_training)
        grads = tape.gradient(g_loss, generator.trainable_variables)
        g_optimizer.apply_gradients(zip(grads, generator.trainable_variables))

        if epoch % 100 == 0:
            print(epoch, 'd-loss:', float(d_loss), 'g-loss:', float(g_loss))
            # 可视化
            z = tf.random.normal([100, z_dim])
            fake_image = generator(z, training=False)
            img_path = os.path.join('gan_images', 'gan-%d.png' % epoch)
            save_result(fake_image.numpy(), 10, img_path, color_mode='P')

        d_losses.append(float(d_loss))
        g_losses.append(float(g_loss))

        if epoch % 10000 == 1:
            generator.save_weights('generator.ckpt')
            discriminator.save_weights('discriminator.ckpt')
Ejemplo n.º 3
0
def main():

    tf.random.set_seed(3333)
    np.random.seed(3333)
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    assert tf.__version__.startswith('2.')

    z_dim = 100  # 隐藏向量z的长度
    epochs = 3000000  # 训练步数
    batch_size = 64  # batch size
    learning_rate = 0.0002
    is_training = True

    # 获取数据集路径
    # C:\Users\z390\Downloads\anime-faces
    # r'C:\Users\z390\Downloads\faces\*.jpg'
    # img_path = glob.glob(r'C:\Users\z390\Downloads\anime-faces\*\*.jpg') + \
    # glob.glob(r'C:\Users\z390\Downloads\anime-faces\*\*.png')
    img_path = glob.glob(
        r'/home/ulysses/workspace/AI/Deep-Learning-with-TensorFlow-book/ch13/faces/*.jpg'
    )
    # img_path.extend(img_path2)
    print('images num:', len(img_path))
    # 构建数据集对象
    dataset, img_shape, _ = make_anime_dataset(img_path, batch_size, resize=64)
    print(dataset, img_shape)
    sample = next(iter(dataset))  # 采样
    print(sample.shape,
          tf.reduce_max(sample).numpy(),
          tf.reduce_min(sample).numpy())
    dataset = dataset.repeat(100)  # 重复循环
    db_iter = iter(dataset)

    generator = Generator()  # 创建生成器
    generator.build(input_shape=(4, z_dim))
    discriminator = Discriminator()  # 创建判别器
    discriminator.build(input_shape=(4, 64, 64, 3))
    # 分别为生成器和判别器创建优化器
    g_optimizer = keras.optimizers.Adam(learning_rate=learning_rate,
                                        beta_1=0.5)
    d_optimizer = keras.optimizers.Adam(learning_rate=learning_rate,
                                        beta_1=0.5)
    if os.path.exists(r'./generator.ckpt.index'):
        generator.load_weights('generator.ckpt')
        print('Loaded generator chpt!!')
    if os.path.exists(r'./discriminator.ckpt.index'):
        discriminator.load_weights('discriminator.ckpt')
        print('Loaded discriminator chpt!!')

    d_losses, g_losses = [], []
    for epoch in range(epochs):  # 训练epochs次
        # 1. 训练判别器
        for _ in range(1):
            # 采样隐藏向量
            batch_z = tf.random.normal([batch_size, z_dim])
            batch_x = next(db_iter)  # 采样真实图片
            # 判别器前向计算
            with tf.GradientTape() as tape:
                d_loss = d_loss_fn(generator, discriminator, batch_z, batch_x,
                                   is_training)
            grads = tape.gradient(d_loss, discriminator.trainable_variables)
            d_optimizer.apply_gradients(
                zip(grads, discriminator.trainable_variables))
        # 2. 训练生成器
        # 采样隐藏向量
        batch_z = tf.random.normal([batch_size, z_dim])
        batch_x = next(db_iter)  # 采样真实图片
        # 生成器前向计算
        with tf.GradientTape() as tape:
            g_loss = g_loss_fn(generator, discriminator, batch_z, is_training)
        grads = tape.gradient(g_loss, generator.trainable_variables)
        g_optimizer.apply_gradients(zip(grads, generator.trainable_variables))

        if epoch % 100 == 0:
            print(epoch, 'd-loss:', float(d_loss), 'g-loss:', float(g_loss))
            # 可视化
            z = tf.random.normal([100, z_dim])
            fake_image = generator(z, training=False)
            img_path = os.path.join('gan_images1', 'gan-%d.png' % epoch)
            save_result(fake_image.numpy(), 10, img_path, color_mode='P')

            d_losses.append(float(d_loss))
            g_losses.append(float(g_loss))

        if epoch % 10000 == 0:
            # print(d_losses)
            # print(g_losses)
            generator.save_weights('./check_point/generator.ckpt')
            discriminator.save_weights('./check_point/discriminator.ckpt')