torchlens
torchlens copied to clipboard
The mode of ech layer
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?
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.
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.