For a previously imported Keras model, after training it with DL4J Spark,
    we want to set the resulting weights back to the original Keras model.

    :param java_model: DL4J model (MultiLayerNetwork or ComputationGraph
    :return: list of numpy arrays in correct order for model.set_weights(...) of a corresponding Keras model

def retrieve_keras_weights(java_model):
    """For a previously imported Keras model, after training it with DL4J Spark,
    we want to set the resulting weights back to the original Keras model.

    :param java_model: DL4J model (MultiLayerNetwork or ComputationGraph
    :return: list of numpy arrays in correct order for model.set_weights(...) of a corresponding Keras model
    """
    weights = []
    layers = java_model.getLayers()
    for layer in layers:
        params = layer.paramTable()
        keys = params.keySet()
        key_list = java_classes.ArrayList(keys)
        for key in key_list:
            weight = params.get(key)
            np_weight = np.squeeze(to_numpy(weight))
            weights.append(np_weight)
    return weights