Модель Tensorflow JS GAN никогда не учится - PullRequest
Новая
       4

Модель Tensorflow JS GAN никогда не учится

1 голос
/ 19 марта 2020

Я пытаюсь перенести пример Python DCGAN MNIST Code Lab (https://www.tensorflow.org/tutorials/generative/dcgan) на Tensorflow. js. Модель генератора должна быть способна создавать изображения рукописных цифр, аналогичные образцам данных MNIST.

Мой код выполняется без ошибок, но я столкнулся с двумя основными проблемами.

  1. Процесс обучения значительно медленнее, чем пример Python. Например, JS в браузере по сравнению с примером Python в Google Code Lab.
  2. Модель My Generator никогда не достигает точки, где она фактически генерирует рукописные числа.

Он учится до такой степени, что генерирует изображения в виде сетки, но, кажется, никогда не узнает многого из этого.

enter image description here

Я считаю, что модели являются Порт 1: 1. Вот мои модели. 

// discriminator model
let dModel = tf.sequential();
const IMAGE_WIDTH = 28;
const IMAGE_HEIGHT = 28;
const IMAGE_CHANNELS = 1;

dModel.add(
tf.layers.conv2d({inputShape: [IMAGE_WIDTH, IMAGE_HEIGHT, IMAGE_CHANNELS], kernelSize: [5,5], filters: 64, strides: [2,2], activation: "relu",
    kernelInitializer: "varianceScaling"
  })
);

dModel.add(tf.layers.leakyReLU())
dModel.add(tf.layers.dropout(0.3))

dModel.add(
tf.layers.conv2d({kernelSize: [5,5], filters: 128, strides: [2,2],
    activation: "relu", kernelInitializer: "varianceScaling"
  })
);

dModel.add(tf.layers.leakyReLU())
dModel.add(tf.layers.dropout(0.3))
dModel.add(tf.layers.flatten());

const NUM_OUTPUT_CLASSES = 1;
dModel.add(tf.layers.dense({units: NUM_OUTPUT_CLASSES}))

// generator model
let gModel = tf.sequential();
gModel.add(tf.layers.dense({units: 7 * 7 * 256,inputShape: [100], useBias: false}));
gModel.add(tf.layers.batchNormalization());
gModel.add(tf.layers.leakyReLU());

gModel.add(tf.layers.reshape({ targetShape: [7, 7, 256] }));

gModel.add(tf.layers.conv2dTranspose({filters: 128, kernelSize: [5, 5], strides: [1, 1], useBias: false, padding: "same"}));
gModel.add(tf.layers.batchNormalization());
gModel.add(tf.layers.leakyReLU());

gModel.add(tf.layers.conv2dTranspose({filters: 64, kernelSize: [5, 5], strides: [2, 2], useBias: false,padding: "same" }));
gModel.add(tf.layers.batchNormalization());
gModel.add(tf.layers.leakyReLU());

gModel.add(tf.layers.conv2dTranspose({filters: 1,kernelSize: [5, 5], strides: [2, 2], useBias: false,padding: "same", activation: "tanh" }));

Функции потерь - это то, где я не могу найти JS, эквивалентный градиентной ленте, поэтому я разработал их немного по-другому.

Пример Python использует: 

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss

def generator_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)

def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
      generated_images = generator(noise, training=True)

      real_output = discriminator(images, training=True)
      fake_output = discriminator(generated_images, training=True)

      gen_loss = generator_loss(fake_output)
      disc_loss = discriminator_loss(real_output, fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

Где, как я использовал optimizer.minimize. Я не уверен, что это может привести к переобучению дискриминатора и вызвать проблему.

Несмотря на то, что это похоже на повторяющиеся вызовы model.predict внутри функций потерь, я сделал это так, иначе я получил ошибку Please make sure the operations that use variables are inside the function f passed to minimize()

function trainStep() {
  const noise = tf.randomNormal([BATCH_SIZE, 100])
  const fakeLabels = tf.ones([BATCH_SIZE], 'int32')
  const realLabels = tf.zeros([BATCH_SIZE], 'int32')

  const dLossCalc = () => {
    const fakeImages = gModel.predict(noise).add(1).div(2)
    let realImages = data.nextTrainBatch(BATCH_SIZE).xs
    realImages = realImages.reshape([BATCH_SIZE, IMAGE_WIDTH, IMAGE_HEIGHT, 1])
    realImages = realImages.sub(127.5).div(127.5)     //normalize to 1, -1

    const fakeLogits = dModel.predict(fakeImages).reshape([BATCH_SIZE])
    const realLogits = dModel.predict(realImages).reshape([BATCH_SIZE])

    const fakeLoss = tf.losses.sigmoidCrossEntropy(fakeLabels.mul(0.98), fakeLogits)
    const realLoss = tf.losses.sigmoidCrossEntropy(realLabels, realLogits)
    const totalLoss = fakeLoss.add(realLoss)
    console.log('Disc Loss ' + totalLoss.dataSync())
    return totalLoss
  }

  const gLossCalc = () => {
    const fakeImages = gModel.predict(noise).add(1).div(2)
    const logits = dModel.predict(fakeImages).reshape([BATCH_SIZE])
    const loss = tf.losses.sigmoidCrossEntropy(fakeLabels, logits) 
    console.log('Gen Loss ' + loss.dataSync())
    return loss
  }

  dOptimizer.minimize(dLossCalc)
  gOptimizer.minimize(gLossCalc)
}

На данный момент я потратил часы и был бы признателен за любую помощь.

Две основные вещи, для которых я не мог найти эквивалент, были градиент Градиенты Tape / Apply и функция потери tf.keras.losses.BinaryCrossentropy. Вместо этого я использую sigmoidCrossEntropy.

Вот полностью рабочий пример codepen, если кто-то хочет посмотреть: https://codepen.io/freeman-g/pen/KKpRyyX?editors=0010

В качестве примечания, я заметил, что applyGradients не задокументировано в документе Tensorflow. js API и открыло связанную проблему GitHub: https://github.com/tensorflow/tfjs/issues/2897

...