Applying TF GradientTape to custom loss function and running training

TensorFlow
, ,
1

Hi,

I have a 3D CNN model that builds a 3D array from random noise. I have a custom loss function that takes in a tensor, converts into numpy(does numpy operations) and returns out a tensor. I am trying to run gradient Tape to compute the loss and gradients and propagate it through the network and I am running into issues. Here is my code

import tensorflow as tf
from tensorflow.keras.layers import *
import numpy as np
import matplotlib.pyplot as plt
import time
import tensorflow.experimental.numpy as tnp
from tensorflow.python.framework.ops import enable_eager_execution
from tensorflow.keras.optimizers import Adam
import math

noise = 200
strides = (2,2,2)
EPOCHS = 25
INIT_LR = 1e-3
BUFFER_SIZE = 1
BATCH_SIZE = 128
opt = Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)

def make_generator_model():
  model = tf.keras.Sequential()
  model.add(tf.keras.Input(shape=(1, 1, 1, noise)))
  model.add(tf.keras.layers.Conv3DTranspose(filters=512, kernel_size=(4,4,4),
                  strides=(1, 1, 1), kernel_initializer='glorot_normal',
                  bias_initializer='zeros', padding='valid'))
  model.add(tf.keras.layers.BatchNormalization())
  model.add(tf.keras.layers.Activation(activation='relu'))

  model.add(tf.keras.layers.Conv3DTranspose(filters=256, kernel_size=(4,4,4),
                  strides=strides, kernel_initializer='glorot_normal',
                  bias_initializer='zeros', padding='same'))
  model.add(tf.keras.layers.BatchNormalization())
  model.add(tf.keras.layers.Activation(activation='relu'))
  model.add(tf.keras.layers.Conv3DTranspose(filters=128, kernel_size=(4,4,4),
                  strides=strides, kernel_initializer='glorot_normal',
                  bias_initializer='zeros', padding='same'))
  model.add(tf.keras.layers.BatchNormalization())
  model.add(tf.keras.layers.Activation(activation='relu'))

  model.add(tf.keras.layers.Conv3DTranspose(filters=64, kernel_size=(4,4,4),
                  strides=strides, kernel_initializer='glorot_normal',
                  bias_initializer='zeros', padding='same'))
  model.add(tf.keras.layers.BatchNormalization())
  model.add(tf.keras.layers.Activation(activation='relu'))
  model.add(tf.keras.layers.Conv3DTranspose(filters=32, kernel_size=(4,4,4),
                  strides=strides, kernel_initializer='glorot_normal',
                  bias_initializer='zeros', padding='same'))
  model.add(tf.keras.layers.BatchNormalization())
  model.add(tf.keras.layers.Activation(activation='sigmoid'))
  model.add(tf.keras.layers.Conv3DTranspose(filters=1, kernel_size=(4,4,4),
                  strides=strides, kernel_initializer='glorot_normal',
                  bias_initializer='zeros', padding='same'))
  return model

def custom_loss(images):
       #Numpy Operations

def step(train_dataset):
  with tf.GradientTape() as gen_tape:
    for x,y in enumerate(train_dataset):
      y = np.expand_dims(y,axis=0)
      generated_images = generator(y, training=True) #Eager Tensor    
      gen_loss = custom_loss(generated_images)

  gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
  generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))


train_images = tf.random.normal([1000, 200])
train_images = tf.expand_dims(train_images, axis=1)
train_images = tf.expand_dims(train_images, axis=2)
train_images = tf.expand_dims(train_images, axis=3)
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE)

generator = make_generator_model()

for epoch in range(0, EPOCHS):
  epochStart = time.time()
  step(train_dataset)

I have the following questions,

  1. Am I implementing this correct?
  2. I sometimes have issue converting ops Tensor to numpy for my loss function
  3. When I run this now, it gives unable to allocate memory. All I am wanting to do is to generate an image from loss