torchlens icon indicating copy to clipboard operation
torchlens copied to clipboard

The mode of ech layer

Open whisperLiang opened this issue 7 months ago • 2 comments

When I forward with each layer in model_history.layer_list, how do I know whether the layer belongs to train mode or eval mode?

whisperLiang avatar May 26 '25 09:05 whisperLiang

Great question—currently there is not a field for this in TorchLens, but there surely should be. I will add this to my to-do list for the next update.

johnmarktaylor91 avatar May 26 '25 12:05 johnmarktaylor91

def train_mode(model):
    """
    Set model to train mode.
    """
    for layer in model:
        # Set the layer to train mode
        if hasattr(layer, "func_argnames") and "training" in layer.func_argnames:
            for idx, arg in layer.func_all_args_non_tensor:
                if isinstance(arg, bool):
                    layer.func_all_args_non_tensor[idx] = True
                    break
    print("train_mode called")
    return model


def eval_mode(model):
    """
    Set model to eval mode.
    """
    for layer in model:
        # Set the layer to eval mode
        if hasattr(layer, "func_argnames") and "training" in layer.func_argnames:
            for idx, arg in layer.func_all_args_non_tensor:
                if isinstance(arg, bool):
                    layer.func_all_args_non_tensor[idx] = False
                    break
    print("eval_mode called")
    return model

Currently, I use this two function to finish it. However, it is not always correct, because the first bool value may not be training value.

whisperLiang avatar May 26 '25 12:05 whisperLiang