Multi-input TF model using TFRecord datasets - tensorflow

Question

I’m trying to create a multi-input single output model in TensorFlow.
I load the data from TFRecs using the get_batched_data fn.
def get_batched_dataset(filenames, batch_size):
dataset = (
tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO)
.map(parse_tfrecord_fn, num_parallel_calls=AUTO)
.map(prepare_sample, num_parallel_calls=AUTO)
.batch(batch_size)
)
return dataset
In the above fn, I do some preprocessing with the loaded data from TFRecs using the prepare_sample fn.
def prepare_sample(features):
image = features['image']
w = tf.shape(image)[0]
h = tf.shape(image)[1]
# some type of preprocessing/data augmentation/transforms
x = {'l_eye': l_eye, 'r_eye':r_eye, 'kps':kps} #l_eye & r_eye are images, kps is numerical data
y = out
return x, y
Below is a very small version of how I’m trying to code my model architecture, just to get an idea.
class cnn_model(layers.Layer):
def __init__(self, name='cnn-model'):
super(cnn_model, self).__init__()
self.conv1 = layers.Conv2D(32, kernel_size=7, strides=2, padding='valid')
self.conv2 = layers.Conv2D(64, kernel_size=5, strides=2, padding='valid')
self.conv3 = layers.Conv2D(128, kernel_size=3, strides=1, padding='valid')
self.bn1 = layers.BatchNormalization(axis = -1, momentum=0.9)
self.bn2 = layers.BatchNormalization(axis = -1, momentum=0.9)
self.bn3 = layers.BatchNormalization(axis = -1, momentum=0.9)
self.leakyrelu = layers.LeakyReLU(alpha=0.01)
self.avgpool = layers.AveragePooling2D(pool_size=2)
self.dropout = layers.Dropout(rate=0.02)
def call(self, input_image):
x = self.conv1(input_image)
x = self.bn1(x)
x = self.leakyrelu(x)
x = self.avgpool(x)
x = self.dropout(x)
x = self.conv2(x)
x = self.bn2(x)
x = self.leakyrelu(x)
x = self.avgpool(x)
x = self.dropout(x)
x = self.conv3(x)
x = self.bn3(x)
x = self.leakyrelu(x)
x = self.avgpool(x)
x = self.dropout(x)
return x
class num_model(layers.Layer):
def __init__(self, name='num-model'):
super(num_model, self).__init__()
self.dense1 = layers.Dense(128)
self.dense2 = layers.Dense(16)
def call(self, input_keypoints):
x = self.dense1(input_keypoints)
x = self.dense2(x)
return x
class main_model(Model):
def __init__(self, name='main-model'):
super(main_model, self).__init__()
self.cnn_model = cnn_model()
self.num_model = num_model()
self.dense1 = layers.Dense(8)
self.dense2 = layers.Dense(2)
def call(self, input_l_r_kps):
leftEye, rightEye, kps = input_l_r_kps['l_eye'], input_l_r_kps['r_eye'], input_l_r_lms['kps']
l_eye_feat = tf.reshape(self.cnn_model(leftEye), (1, 3*128*128))
r_eye_feat = tf.reshape(self.cnn_model(rightEye), (1, 3*128*128))
kp_feat = self.num_model(kps)
combined_feat = tf.concat((l_eye_feat, r_eye_feat, lm_feat),1)
x = self.dense1(combined_feat)
x = self.dense2(x)
return x
Now, the dataset returned by the get_batched_dataset fn is what I’ll be feeding into the Keras model.fit method.
train_dataset = get_batched_dataset('train.tfrec', batch_size)
valid_dataset = get_batched_dataset('valid.tfrec', batch_size)
model.fit(
x=train_dataset,
batch_size=batch_size,
epochs=1,
validation_data=valid_dataset,
use_multiprocessing=False
)
Can you please guide me where I’m going wrong? Is it in the prepare_sample fn by returning x as a dict, or somewhere in the model code? I’m really new to TF and confused.
Any help appreciated!