[FEATURE] how to return mlx intermediate layer output similarly to Keras

Describe the FEATURE

Is there a way to get the intermediate tensor results similar to keras following code. I try to determine why the mlx/ Keras code are not returning the same output. This is almost one week that I work on it and when I apply keras pretrained values I got a completely different result. I would like to determine where the drift occur and fix it without making N intermediates models and reuse the final Keras model weights.

To Reproduce

Include code snippet


from tensorflow.keras.models import Model

def get_intermediate_outputs_keras(model, input_data):
    layer_outputs = [layer.output for layer in model.layers]  # all layer outputs
    intermediate_model = Model(inputs=model.input, outputs=layer_outputs)
    intermediate_predictions = intermediate_model.predict(input_data)
    return intermediate_predictions

Expected behavior A clear and concise description of what you expected to happen.

Desktop (please complete the following information):

OS Version: MacOS 14.4
Version 0.7.0

Additional context Add any other context about the problem here.

There are two things to address here, one debugging the outputs of the models and the other is capturing intermediate values from a model.

For the 1st I usually just put a debugger and step layer by layer inspecting the intermediates directly in the debugger. This is much better as it has infinite granularity compared to capturing some specific intermediates. You can probably achieve the same in Keras maybe with the PyTorch backend? Not sure.

For the 2nd there isn't a generic way to capture intermediates. You can simply store them in a global container if it is for debugging purposes or even make a wrapper that does it for you something like the following.

class DebugWrapper(nn.Module):
    outputs = []

    def __init__(self, wrapped):
        super().__init__()
        self.wrapped = wrapped

    def __call__(self, *args, **kwargs):
        out = self.wrapped(*args, **kwargs)
        self.outputs.append(out)
        return out

    @classmethod
    def clear_outputs(cls):
        cls.outputs.clear()

model = ...
for idx in range(len(model.layers)):
    model.layers[idx] = DebugWrapper(model.layers[idx])

In general, when it is possible to have various outputs from a model, a custom output data class is used that may be partially filled depending on the config. I think this has been found to be slightly less brittle and more descriptive than hooks. Otoh Keras' intermediate outputs and losses is one of my favorite features that it has.

Since I don't think that we will provide a specific way to capture intermediates, let us know if this answer covers your use case and then we can close the issue.

Yes your proposal is interesting. That would be nice to have access to the layers in order to run this DebugWrapper class. In keras this is of course easy.

Le mar. 30 avr. 2024 à 20:19, Angelos Katharopoulos < @.***> a écrit :

There are two things to address here, one debugging the outputs of the models and the other is capturing intermediate values from a model.

For the 1st I usually just put a debugger and step layer by layer inspecting the intermediates directly in the debugger. This is much better as it has infinite granularity compared to capturing some specific intermediates. You can probably achieve the same in Keras maybe with the PyTorch backend? Not sure.

For the 2nd there isn't a generic way to capture intermediates. You can simply store them in a global container if it is for debugging purposes or even make a wrapper that does it for you something like the following.

class DebugWrapper(nn.Module): outputs = []
def __init__(self, wrapped):
    super().__init__()
    self.wrapped = wrapped

def __call__(self, *args, **kwargs):
    out = self.wrapped(*args, **kwargs)
    self.outputs.append(out)
    return out

@classmethod
def clear_outputs(cls):
    cls.outputs.clear()
model = ...for idx in range(len(model.layers)): model.layers[idx] = DebugWrapper(model.layers[idx])

In general, when it is possible to have various outputs from a model, a custom output data class is used that may be partially filled depending on the config. I think this has been found to be slightly less brittle and more descriptive than hooks. Otoh Keras' intermediate outputs and losses is one of my favorite features that it has.

Since I don't think that we will provide a specific way to capture intermediates, let us know if this answer covers your use case and then we can close the issue.

— Reply to this email directly, view it on GitHub https://github.com/ml-explore/mlx/issues/1056#issuecomment-2086402039, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJBWYSCHGPEKLYYJSZZD23Y77OCHAVCNFSM6AAAAABHAAZKISVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAOBWGQYDEMBTHE . You are receiving this because you authored the thread.Message ID: @.***>

ml-explore / mlx

[FEATURE] how to return mlx intermediate layer output similarly to Keras #1056