dlr eo: set custom_objects when loading model

This commit is contained in:
Starbeamrainbowlabs 2023-03-01 17:19:10 +00:00
parent b5f23e76d1
commit f70083bea4
Signed by: sbrl
GPG key ID: 1BE5172E637709C2

View file

@ -156,7 +156,10 @@ if PATH_CHECKPOINT is None:
model = DeeplabV3Plus(image_size=IMAGE_SIZE, num_classes=NUM_CLASSES, num_channels=8) model = DeeplabV3Plus(image_size=IMAGE_SIZE, num_classes=NUM_CLASSES, num_channels=8)
summarywriter(model, os.path.join(DIR_OUTPUT, "summary.txt")) summarywriter(model, os.path.join(DIR_OUTPUT, "summary.txt"))
else: else:
model = tf.keras.models.load_model(PATH_CHECKPOINT) model = tf.keras.models.load_model(PATH_CHECKPOINT, custom_objects={
# Tell Tensorflow about our custom layers so that it can deserialise models that use them
"LossCrossEntropyDice": LossCrossEntropyDice
})