keras 无法识别自定义对象

Question

I am loading my model using custom a custom loss function but when I run the code I get an error: ValueError: Unknown loss function:dice_coef_loss . The was created using 2 GPUs. When I save the model using 1 GPU the load_model() I don't get the error. Is there a reason why a multi-gpu trained model would not recognize custom_objects ?

import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras import backend as K

def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)

load_model('test_2gpus_model', custom_objects = {'dice_coef': dice_coef, 'dice_coef_loss': dice_coef_loss}))

I have also tried:

load_model('test_2gpus_model', custom_objects = {'dice_coef': dice_coef(y_true, y_pred), 'dice_coef_loss': dice_coef_loss(y_true, y_pred)}))

but then I get the error NameError: name 'y_true' is not defined

Answer 1

I had the same problem because the function definition was in another class other than the one I am calling the load_model from. when I copied the function definitions inside the class that calls the load_model, it worked.