How to use a pre-trained keras model for inference in tf.data.Dataset.map?

Question

I have a pre-trained model, and I'm trying to build another model that takes as input the output of the previous model. I don't want to train the models end-to-end, and want to use the first model for inference only. The first model was trained using tf.data.Dataset pipeline, and my first inclination was to integrate the model as just another dataset.map() operation at the tail of the pipeline, but I'm having issues with that. I've encountered 20 different errors in the process, each unrelated to the prior one. The batch-normalization layer particularly seems to be a pain-point for this.

Below is a minimal starter example that illustrates the issue. It's written in R but answers in python are also welcome.

I'm using tensorflow-gpu version 1.13.1 and keras from tf.keras

library(reticulate)
library(tensorflow)
library(keras)
library(tfdatasets)
use_implementation("tensorflow")

model_weights_path <- 'model-weights.h5'

arr <- function(...) 
  np_array(array(seq_len(prod(unlist(c(...)))), unlist(c(...))), dtype = 'float32')

new_model <- function(load_weights = TRUE) {
  model <- keras_model_sequential() %>% 
    layer_conv_1d(5, 5, activation = 'relu', input_shape = shape(150, 10)) %>%
    layer_batch_normalization() %>%
    layer_flatten() %>%
    layer_dense(10, activation = 'softmax')
  if (load_weights)
    load_model_weights_hdf5(model, model_weights_path)
  freeze_weights(model)
  model
}

if(!file.exists(model_weights_path)) {
  model <- new_model(FALSE) 
  save_model_weights_hdf5(model, model_weights_path)
}

model <- new_model()

data <- arr(20, 150, 10)
ds <- tfdatasets::tensors_dataset(data) %>% 
  dataset_repeat()

ds2 <- ds %>% 
  dataset_map(function(x) {
    model(x)
  })

try(nb <- next_batch(ds2))

sess <- k_get_session()
it <- make_iterator_initializable(ds2)
sess$run(iterator_initializer(it))
nb <- it$get_next()

try(sess$run(nb))

sess$run(tf$initialize_all_variables())

try(sess$run(nb))

Answer 1

Maybe this will not answer your question directly as I am not familiar with R. But I have recently built an input pipeline using tf.data .

The generate_images function is mapped using .map and uses a trained generator model to generate new images.

gen_model = tf.keras.models.load_model(artifact_dir+'/'+generators[-1], compile=False)

NOISE_DIM = 100

def generate_images(l):
    # generate images using the trained generator
    noise = tf.random.normal([BATCH_SIZE, NOISE_DIM])
    images = gen_model(noise)

    # prepare the images for resize_and_preprocess function
    images = tf.squeeze(images, axis=-1)
    images = images*0.5+0.5
    images = tf.image.convert_image_dtype(images, dtype=tf.uint8)

    return images

genloader = tf.data.Dataset.from_tensors([1])

genloader = (
    genloader
    .map(generate_images, num_parallel_calls=AUTO)
    .map(resize_and_preprocess, num_parallel_calls=AUTO)
    .prefetch(AUTO)
)

Regarding Batch Normalization, it behaves differently during the training and inference phase. In Python-based TensorFlow one needs to pass training=False when using a pretrained model that has batch normalization layer.