TensorFlow training gets slower every batch - performance

Question

I am new to TensorFlow and I get my code running successfully by modifying tutorials from the official website.
I checked some other answers on StackOverflow, which says my problem is likely due to something is being added to the graph every time. However, I have no idea where to look for the code that might have caused this.
Also, I used tf.py_function to map the dataset because I really need to enable eagerly mode in the mapping.
def get_dataset(data_index):
# data_index is a Pandas Dataframe that contains image/label pair info, each row is one pair
data_index = prepare_data_index(data_index)
# shuffle dataframe here because dataset.shuffle is taking very long time.
data_index = data_index.sample(data_index.shape[0])
path = path_to_img_dir
# list of dataframe indices indicating rows that are going to be included in the dataset for training.
indices_ls = ['{}_L'.format(x) for x in list(data_index.index)] + ['{}_R'.format(x) for x in list(data_index.index)]
# around 310k images
image_count = len(indices_ls)
list_ds = tf.data.Dataset.from_tensor_slices(indices_ls)
# dataset.shuffle is commented out because it takes too much time
# list_ds = list_ds.shuffle(image_count, reshuffle_each_iteration=False)
val_size = int(image_count * 0.2)
train_ds = list_ds.skip(val_size)
val_ds = list_ds.take(val_size)
def get_label(index):
index = str(np.array(index).astype(str))
delim = index.split('_')
state = delim[1]
index = int(delim[0])
if state == 'R':
label = data_index.loc[index][right_labels].to_numpy().flatten()
elif state == 'L':
label = data_index.loc[index][left_labels].to_numpy().flatten()
return tf.convert_to_tensor(label , dtype=tf.float16)
def get_img(index):
index = str(np.array(index).astype(str))
delim = index.split('_')
state = delim[1]
index = int(delim[0])
file_path = '{}_{}.jpg'.format(data_index.loc[index, 'sub_folder'],
str(int(data_index.loc[index, 'img_index'])).zfill(4)
)
img = tf.io.read_file(os.path.join(path, file_path))
img = tf.image.decode_jpeg(img, channels=3)
full_width = 320
img = tf.image.resize(img, [height, full_width])
# Crop half of the image depending on the state
if state == 'R':
img = tf.image.crop_to_bounding_box(img, offset_height=0, offset_width=0, target_height=height,
target_width=int(full_width / 2))
img = tf.image.flip_left_right(img)
elif state == 'L':
img = tf.image.crop_to_bounding_box(img, offset_height=0, offset_width=int(full_width / 2), target_height=height,
target_width=int(full_width / 2))
img = tf.image.resize(img, [height, width])
img = tf.keras.preprocessing.image.array_to_img(
img.numpy(), data_format=None, scale=True, dtype=None
)
# Apply auto white balancing, output an np array
img = AWB(img)
img = tf.convert_to_tensor(img, dtype=tf.float16)
return img
def process_path(index):
label = get_label(index)
img = get_img(index)
return img, label
AUTOTUNE = tf.data.experimental.AUTOTUNE
train_ds = train_ds.map(lambda x: tf.py_function(
process_path,
[x], (tf.float16, tf.float16)), num_parallel_calls=AUTOTUNE)
val_ds = val_ds.map(lambda x: tf.py_function(
process_path,
[x], (tf.float16, tf.float16)), num_parallel_calls=AUTOTUNE)
def configure_for_performance(ds):
ds = ds.cache()
# ds = ds.shuffle(buffer_size=image_count)
ds = ds.batch(batch_size)
ds = ds.prefetch(buffer_size=AUTOTUNE)
return ds
train_ds = configure_for_performance(train_ds)
val_ds = configure_for_performance(val_ds)
return train_ds, val_ds
Can anyone please help me? Thanks!
Here is the rest of my code.
def initialize_model():
IMG_SIZE = (height, width)
preprocess_input = tf.keras.applications.vgg19.preprocess_input
IMG_SHAPE = IMG_SIZE + (3,)
base_model = tf.keras.applications.VGG19(input_shape=IMG_SHAPE,
include_top=False,
weights='imagenet')
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
prediction_layer = tf.keras.layers.Dense(class_num, activation=tf.nn.sigmoid, use_bias=True)
inputs = tf.keras.Input(shape=(height, width, 3))
x = preprocess_input(inputs)
x = base_model(x, training=True)
x = global_average_layer(x)
outputs = prediction_layer(x)
model = tf.keras.Model(inputs, outputs)
def custom_loss(y_gt, y_pred):
L1_loss_out = tf.math.abs(tf.math.subtract(y_gt, y_pred))
scaler = tf.pow(50.0, y_gt)
scaled_loss = tf.math.multiply(L1_loss_out, scaler)
scaled_loss = tf.math.reduce_mean(
scaled_loss, axis=None, keepdims=False, name=None
)
return scaled_loss
base_learning_rate = 0.001
model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=base_learning_rate, momentum=0.9),
loss=custom_loss,
metrics=['mean_absolute_error']
)
return model
def train(data_index, epoch_num, save_path):
train_dataset, validation_dataset = get_dataset(data_index)
model = initialize_model()
model.summary()
history = model.fit(train_dataset,
epochs=epoch_num,
validation_data=validation_dataset)
model.save_weights(save_path)
return model, history