Saving Models in Active Learning setting

Some context, currently I am running an active learning loop. Initially I enable dropouts for my bert model while making predictions to gather uncertainties and to select the most uncertain samples. Now after doing that I fine tune my model based on the initial weights from HF along with new samples I have gathered from the active learning loop. The AL loop looks something like the following:

  1. Fine tune model
  2. Enable dropouts while making predictions(MCdropout) and gather uncertainties
  3. Save Model if needed
  4. Select most informative and uncertain points based on point 2 and add to train and remove from pool
  5. Initialize the model weights from the actual HF model and go to step 1 to run fine tune again and so on

Now in step 3 when I try to save model as trainer.save_model('path/to/model_dir) does that ensure that my model is being saved in a model.eval() state. My assumption was that when I save the model it would be saved in the train state with dropouts enabled(since I had manually enabled dropouts and never switched them off) but that is not the case. When I try to load the model it is no longer in training state(Which is essentially what I want, but I want to make sure that this is the expected behavior out of trainer.save() or is something off)

from transformers import AutoModelForSequenceClassification
from transformers import TextClassificationPipeline
import os

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)

# Fine tune model
trainer.train()

# Explicityly Enable dropouts in prediction(MCdropuout)
for m in model.modules():
  if m.__class__.__name__.startswith('Dropout'):
    m.train()

print(f"Droputs enabled: {model.dropout.training}") # True here

# Save model
trainer.save_model(os.path.join(os.getcwd(),'model'))

# Suppose now you load this model at some later point in time
loaded_model = AutoModelForSequenceClassification.from_pretrained('model')

print(f"Droputs enabled: {loaded_model.dropout.training}") # False here, which is perfect. But is this expected?

The logs are much longer and this is not the actual code but the above code is the general idea. It is fine tuning docs from HF Fine-tune a pretrained model.

You might be tempted to conclude that I can just save my model before I enable dropouts for predictions, but think of this running in a loop. At the nth loop where n!=0(since active learning would not work best at your first pass), dropouts are already enabled for predictions.

To be more precise

  • Looking for inputs on the behavior of trainer.save() in the above scenario. While saving the model with all dropouts enabled what would/should trainer.save() do?

  • Based on the above question, do you suggest that I must disable dropouts and then save the model ? If yes then how, since I only know about trainer.save() which will save my trained model. Please suggest if there are any alternative methods to save

Some References I have already checked

Any ideas for this question? :eyes:

Please let me know if I need to add more context/clarity to this question.