Aleph-Alpha/magma

Subsequent inference calls produce less good results

Closed this issue · 1 comments

Following the code in README.md or example_inference.py to perform inference by calling model.preprocess_inputs(…) followed by model.generate(…) produces good results the first time the pair is called, but poor results for subsequent pairs of calls.

The reason is that model = Magma.from_checkpoint(…) loads the model with inconsistent training/eval settings. model.training is True but model.image_prefix.enc.training is False. The first call to model.preprocess_inputs(…) works correctly as the image encoder has training False and so its Batch Normalisation steps work correctly. The call to model.generate(…) records the training state on entry and restores it on exit, which because model.training is True puts the whole model into training state. Subsequent calls to model.preprocess_inputs(…) then don't perform Batch Normalisation steps correctly.

The play space at https://huggingface.co/spaces/EleutherAI/magma has this problem too.

The fix is to add model.eval() after model = Magma.from_checkpoint(…), setting the whole model to a consistent eval state.

Hi @steve-barlow 👋
Thanks for pointing out this issue -- I've added an eval() here which should fix the problem 🙂