Search code examples
pythontensorflowimage-segmentationtensorflow-datasetstfrecord

How to feed tfrecord file in a model and train?


I wrote a tfrecord file and fed in my Unet model but got a problem with the input shape. Below is my code.

About the data:

  • 484 training images, each has a shape of (240, 240, 155, 4), these 4 numbers are the height, width, number of layers and channels respectively.
  • 484 labels, each has a shape of (240, 240, 155)

I used the first 2 examples:

test_writer = tf.io.TFRecordWriter('test.tfrecords')

for i in range(2):
  example = create_example(image_paths[i], label_paths[i])
  test_writer.write(example.SerializeToString())

test_writer.close()

serialised_example = tf.data.TFRecordDataset('test.tfrecords')
parsed_example = serialised_example.map(parse_tfrecord)

My model architecture (I simplified it):

from tensorflow.keras.layers import Conv3D, Conv3DTranspose, Input, Rescaling

num_classes = 4

my_model = tf.keras.Sequential([

Input(shape = (240, 240, 155, 4)),
Rescaling(scale = 1./255),

Conv3D(filters = 64, kernel_size = 3, strides = 2, activation = 'relu', padding = 'same'),
Conv3D(filters = 64, kernel_size = 3, activation = 'relu', padding = 'same'),
# and more layers between...
Conv3DTranspose(filters = 64, kernel_size = 3, activation = 'relu', padding = 'same'),
Conv3DTranspose(filters = 64, kernel_size = 3, strides = 2, activation = 'relu', padding = 'same'),

Conv3D(filters = num_classes, kernel_size = 3, activation = 'softmax', padding = 'same')

])

my_model.compile(optimizer = 'rmsprop', loss = 'sparse_categorical_crossentropy')

I got my data set from my tfrecord file like this:

def get_image_and_label(features):
  image, label = features['image'], features['label']
  return image, label

def get_dataset(tfrecord_names):

  dataset = (tf.data.TFRecordDataset(tfrecord_names)
             .map(parse_tfrecord)
             .map(get_image_and_label))

  return dataset

dataset = get_dataset('test.tfrecords')

I started to train:

my_model.fit(dataset, epochs = 1)

And got this error: Input 0 of layer "sequential_2" is incompatible with the layer: expected shape=(None, 240, 240, 155, 4), found shape=(240, 240, 155, 4)

How can I fix this? Please tell me if you need more information (data link or my previous code).


Solution

  • Your model expects the shape (samples, 240, 240, 155, 4), so try something like this:

    dataset = get_dataset('test.tfrecords').batch(1)
    

    And you have to set strides=1 if you want your labels to match your output:

    from tensorflow.keras.layers import Conv3D, Conv3DTranspose, Input, Rescaling
    
    num_classes = 4
    
    my_model = tf.keras.Sequential([
    
    Input(shape = (240, 240, 155, 4)),
    Rescaling(scale = 1./255),
    Conv3D(filters = 64, kernel_size = 3, strides = 1, activation = 'relu', padding = 'same'),
    Conv3D(filters = 64, kernel_size = 3, activation = 'relu', padding = 'same'),
    # and more layers between...
    Conv3DTranspose(filters = 64, kernel_size = 3, activation = 'relu', padding = 'same'),
    Conv3DTranspose(filters = 64, kernel_size = 3, strides = 1, activation = 'relu', padding = 'same'),
    
    Conv3D(filters = num_classes, kernel_size = 3, activation = 'softmax', padding = 'same')
    ])