How to log custom Pytorch model with Mlflow?

Question

I have been using Presumm https://github.com/nlpyang/PreSumm for text summarization.

However, in src/train_abstractive.py, the model learner trainer is not a torch.nn.Module. However, the input AbsSummarizer is an extension of the torch.nn.Module class.

I want to use mlflow.pytorch.log_model to save the model as a native pytorch model. But trainer is not a torch.nn.Module. How do I go about this?

Answer 1

After trainer runs, you should log the AbsSummarizer, which should be populated with the updated weights that correspond to the trained model.