import tensorflow as tf

  train = dataset['train'].map(
      load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  test = dataset['test'].map(load_image_test)

  train_dataset = train.cache().shuffle(1000).batch(batch_size).repeat()
  train_dataset = train_dataset.prefetch(
      buffer_size=tf.data.experimental.AUTOTUNE)
  test_dataset = test.batch(batch_size)
  config = hparams_config.get_efficientdet_config('efficientdet-d0')
  config.heads = ['segmentation']
  model = efficientdet_keras.EfficientDetNet(config=config)
  model.build((1, 512, 512, 3))
  model.compile(
      optimizer='adam',
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=['accuracy'])

  val_subsplits = 5
  val_steps = info.splits['test'].num_examples // batch_size // val_subsplits
  model.fit(
      train_dataset,
      epochs=20,
      steps_per_epoch=steps_per_epoch,
      validation_steps=val_steps,
      validation_data=test_dataset,
      callbacks=[])

  model.save_weights('./test/segmentation')
