What are the PyTorch's model.eval() + no_grad() equivalent in TensorFlow?

user_19391408 0
Viewed 247

I am trying to extract BERT embeddings and reproduce this code using tensorflow instead of pytorch. I know tf.stop_gradient() is the equivalent of torch.no_grad() but what about model.eval() / combination of both ?

# Put the model in "evaluation" mode, meaning feed-forward operation.
model.eval()

# Run the text through BERT, and collect all of the hidden states produced
# from all 12 layers. 
with torch.no_grad():

    outputs = model(tokens_tensor, segments_tensors)

    # Evaluating the model will return a different number of objects based on 
    # how it's  configured in the `from_pretrained` call earlier. In this case, 
    # becase we set `output_hidden_states = True`, the third item will be the 
    # hidden states from all layers. See the documentation for more details:
    # https://huggingface.co/transformers/model_doc/bert.html#bertmodel
    hidden_states = outputs[2]
pytorch python
1 Answers

TLDR; eval and no_grad are two completely different things but will often be used in conjunction, primarily for performing fast inference in the case of evaluation/testing loops.

The nn.Module.eval function is applied on a PyTorch module and gives it the ability to change its behaviour depending on the stage type: training or evaluation. Only for a handful of layers does this actually have an effect on the layer. Functions such as dropout layers and normalization layers have different behaviours depending on whether they are in training or evaluation mode. You can read more about it on this thread.

The torch.no_grad utility is however a context manager, it changes the way the code contained inside that scope runs. When applied no_grad has the effect of preventing gradient computation. In practice, this means no layer activation is been cached in memory. This is most generally used for evaluation and testing loops where no backpropagation is expected after an inference. However, it can also be used during training, for example when an inference on a frozen component and the gradient is not required to pass through it.

user_63313690
Related