Training Style Transfer Neural Network Leads to NANs

Question

I've been encountering NANs while trying to train the Neural Style Transfer model from https://arxiv.org/pdf/1610.07629.pdf. I've tried reducing learning rate and using different initialisations but it doesn't seem to work. I suspect it has something to do with my training loop using tape.Gradient.

Here is my Image Transform Network:

import tensorflow as tf
from tensorflow.keras.layers import Conv2D, Input, BatchNormalization, Add, ReLU, Reshape, UpSampling2D
from tensorflow.keras.models import Model
import numpy as np
from tensorflow.keras import initializers

# initializers
initializer = tf.random_normal_initializer(mean=0.0, stddev=0.01, seed=None)
betaInitializer = initializers.constant(0.)
gammaInitializer = initializers.constant(1.)

# Conditional Instance Normalisation layer
class ConditionalInstanceNorm(tf.keras.layers.Layer):
  def __init__(self, scope_bn, y1, y2, alpha):
    super(ConditionalInstanceNorm, self).__init__()
    self.scope_bn = scope_bn
    self.y1 = y1
    self.y2 = y2
    self.alpha = alpha
  
  def build(self, input_shape):
    self.beta = self.add_weight(name="beta"+self.scope_bn, shape=(self.y1.shape[-1], input_shape[-1]), initializer=betaInitializer, trainable=True)
    self.gamma = self.add_weight(name="gamma"+self.scope_bn, shape=(self.y1.shape[-1], input_shape[-1]), initializer=gammaInitializer, trainable=True)
  
  def call(self, inputs):
    mean, var = tf.nn.moments(x=inputs, axes=[1,2], keepdims=True)
    beta1 = tf.matmul(self.y1, self.beta)
    gamma1 = tf.matmul(self.y1, self.gamma)
    beta2 = tf.matmul(self.y2, self.beta)
    gamma2 = tf.matmul(self.y2, self.gamma)
    beta = self.alpha*beta1 + (1. - self.alpha)*beta2
    gamma = self.alpha*gamma1 + (1. - self.alpha)*gamma2
    x = tf.nn.batch_normalization(x=inputs, mean=mean, variance=var, offset=beta, scale=gamma, variance_epsilon=1e-10)
    return x

# Applies upsampling if stride = 0.5, includes mirror padding, conv layer and conditional instance Norm layer
def PadConvBatch(x, filters=32, kernel_size=3, strides=1, activation='relu', scope_bn="", y1=None, y2=None, alpha=1):
  if isinstance(strides, float):
    x = UpSampling2D(size=2, interpolation='nearest')(x)
    strides=1
  padding = tf.cast((kernel_size-1)/2, tf.int32)
  x = tf.pad(x, [[0,0], [padding,padding], [padding,padding], [0,0]], "REFLECT")
  x = Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, activation=activation, kernel_initializer=initializer)(x)
  conditionalInstanceNorm = ConditionalInstanceNorm(scope_bn=scope_bn, y1=y1, y2=y2, alpha=alpha)
  x = conditionalInstanceNorm(x)
  return x

# Implementation of resnet block according to paper
def resblock(x, scope_bn1="", scope_bn2="", y1=None, y2=None, alpha=1):
  fx = PadConvBatch(x=x, filters=128, activation="linear", scope_bn=scope_bn1, y1=y1, y2=y2, alpha=alpha)
  fx = ReLU()(fx)
  fx = PadConvBatch(x=fx, filters=128, activation="linear", scope_bn=scope_bn2, y1=y1, y2=y2, alpha=alpha)
  out = Add()([x,fx])
  return out

# Overall network
def ImageTransformNetwork(shape=(256,256,3), y1=None, y2=None, alpha=1):
  inputs = Input(shape=shape) #1
  x = PadConvBatch(x=inputs, kernel_size=9, scope_bn="1", y1=y1, y2=y2, alpha=alpha) #3
  x = PadConvBatch(x=x, filters=64, strides=2, scope_bn="2", y1=y1, y2=y2, alpha=alpha) #5
  x = PadConvBatch(x=x, filters=128, strides=2, scope_bn="3", y1=y1, y2=y2, alpha=alpha) #7
  x = resblock(x=x, scope_bn1="4", scope_bn2="5", y1=y1, y2=y2, alpha=alpha) #13
  x = resblock(x=x, scope_bn1="6", scope_bn2="7", y1=y1, y2=y2, alpha=alpha) #19
  x = resblock(x=x, scope_bn1="8", scope_bn2="9", y1=y1, y2=y2, alpha=alpha) #25
  x = resblock(x=x, scope_bn1="10", scope_bn2="11", y1=y1, y2=y2, alpha=alpha) #31
  x = resblock(x=x, scope_bn1="12", scope_bn2="13", y1=y1, y2=y2, alpha=alpha) #37
  x = PadConvBatch(x=x, filters=64, strides=0.5, scope_bn="14", y1=y1, y2=y2, alpha=alpha) #39
  x = PadConvBatch(x=x, filters=32, strides=0.5, scope_bn="15", y1=y1, y2=y2, alpha=alpha) #41
  x = tf.pad(x, [[0,0], [4,4], [4,4], [0,0]], "REFLECT")
  x = Conv2D(filters=3, kernel_size=9, strides=1, activation='sigmoid', kernel_initializer=initializer)(x) #42
  x = x*255
  model = Model(inputs=inputs, outputs=x)
  return model

Here is my loss network:

import tensorflow as tf
import os

#config
os.environ['TFHUB_MODEL_LOAD_FORMAT'] = 'COMPRESSED'

#Specify layers for content and style representation
content_layers = ['block4_conv2']
style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']

# Find number of layers
num_content_layers = len(content_layers)
num_style_layers = len(style_layers)

# Grab vgg19 network without dense layers
def vgg_layers(layer_names):
  vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')

  outputs = [vgg.get_layer(name).output for name in layer_names]

  model = tf.keras.Model([vgg.input], outputs)
  return model

# Calculate normalised gram matrix
def gram_matrix(input_tensor):
  result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
  input_shape = tf.shape(input_tensor)
  num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32)
  return result/(num_locations)

# Setting up loss network model class + forward method
class LossNetwork(tf.keras.models.Model):
  def __init__(self, style_layers, content_layers):
    super(LossNetwork, self).__init__()
    self.vgg = vgg_layers(style_layers+content_layers)
    self.style_layers = style_layers
    self.content_layers = content_layers
    self.num_style_layers = len(style_layers)
    self.vgg.trainable = False
  
  def call(self, inputs):
    inputs = inputs*255.0
    preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs)
    outputs = self.vgg(preprocessed_input)
    style_outputs, content_outputs = (outputs[:self.num_style_layers], outputs[self.num_style_layers:])

    style_outputs = [gram_matrix(style_output) for style_output in style_outputs]
    content_dict = {content_name: value for content_name, value in zip(self.content_layers, content_outputs)}
    style_dict = {style_name: value for style_name, value in zip(self.style_layers, style_outputs)}
    return {'content': content_dict, 'style': style_dict}

# returns network
def loadLossNetwork():
  return LossNetwork(style_layers, content_layers)

As for training, I'm currently testing it out by only using batch size of 1 and using only 1 image. Here is my training code:

import tensorflow as tf
import numpy as np
import time
import functools

import ImageTransformNetwork
import ImagePreProcessing
import LossNetwork

import PIL.Image

"""Utility and Loss Functions START"""
#Converts a tensor into an image
def tensor_to_image(tensor):
  #Multiplies every element in tensor by 255
  tensor = tensor*255

  #Converts tensor into a numpy array
  tensor = np.array(tensor, dtype=np.uint8)

  if np.ndim(tensor)>3:
    #Ensures only 1 image is being sent
    assert tensor.shape[0] == 1

    #Takes out the first element so you just get a 3 dim np array
    tensor = tensor[0]
  
  #Converts numpy array into PIL image
  return PIL.Image.fromarray(tensor)

# Calculates style and content loss in loss network
def style_content_loss(outputs):
  style_outputs = outputs['style']
  content_outputs = outputs['content']

  style_loss_total = tf.constant(0.0, tf.float32)
  for name in style_outputs.keys():
    style_loss_total += tf.reduce_sum((style_outputs[name]-style_targets[name])**2)*style_weight/batch_size
  
  content_loss_total = 0
  for name in content_outputs.keys():
    content_loss_total += ((content_outputs[name] - content_targets[name])**2)
    content_loss_total *= content_weight/(content_outputs[name].shape[1]*content_outputs[name].shape[2])

  loss = style_loss_total + content_loss_total

  return loss

"""Utility and Loss Functions END"""


"""Setup START"""
#Loads images
content_image, style_image = ImagePreProcessing.loadImage()
content_image = tf.image.resize(content_image, [256,256])
content_image = tf.reshape(content_image, [1,256,256,3])

y1 = tf.Variable(tf.zeros([1,10]))
y2 = tf.Variable(tf.zeros([1,10]))
alpha = tf.constant([1.])

#Calls instance of image transform network and loss network
imageTransformNetwork = ImageTransformNetwork.ImageTransformNetwork(y1=y1, y2=y2, alpha=alpha)
lossNetwork = LossNetwork.loadLossNetwork()

# Obtain features of each layer of loss network model
style_targets = lossNetwork(style_image)['style']
content_targets = lossNetwork(content_image)['content']

#Sets optimiser, LBFGS is better
opt = tf.keras.optimizers.Adam(learning_rate=1e-3, beta_1=0.9, beta_2=0.999, epsilon=1e-1)

"""Setup END"""


"""Train START"""

#Training step
style_weight = 1e-2
content_weight = 1e4
total_variation_weight = 30

@tf.function()
def train_step(image):
  with tf.GradientTape() as tape:
    stylised = imageTransformNetwork(image, training=True)
    outputs = lossNetwork(stylised)

    loss = style_content_loss(outputs)
  
  grad = tape.gradient(loss, imageTransformNetwork.trainable_variables)
  opt.apply_gradients(zip(grad, imageTransformNetwork.trainable_variables))

#Optimisation loop
epochs = 1
steps_per_epoch = 5
batch_size = 1

start = time.time()
step = 0
for n in range(epochs):
  for m in range(steps_per_epoch):
    step += 1
    train_step(content_image)
    print(".", end='', flush=True)
  print('Train step: {}'.format(step))
end=time.time()
print('Total time: {:.1f}'.format(end-start))

print(imageTransformNetwork.layers[15].weights)
stylised_tensor = imageTransformNetwork(content_image)
print(stylised_tensor)
stylised_image = tensor_to_image(stylised_tensor)
stylised_image.show()

"""Train END"""

Here are the weights of my Action Transform Network after 5 epochs:

[

    [[nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     ...,
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan]],

    [[nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     ...,
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan]]],


   [[[nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     ...,
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan]],

    [[nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     ...,
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan]],

    [[nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     ...,
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan]]],


   [[[nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     ...,
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan]],

    [[nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     ...,
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan]],

    [[nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     ...,
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan],
     [nan, nan, nan, ..., nan, nan, nan]]]], dtype=float32)>,


array([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
dtype=float32)>]


Is there anything wrong with my code that could lead to this?

Training Style Transfer Neural Network Leads to NANs

Answers (1)

Related Questions