Search code examples
pythonneural-networkpysparkpredictionautoencoder

Saving and loading an autoencoder

I am trying to make a predictor using Auto Encoder (AE) and Alternating Least Squares (ALS) methods. ALS is created using pyspark.mllib.recommendation package. I am able to save the ALS model and reuse it by model.save() and pyspark.mllib.recommendation.MatrixFactorizationModel.load() methods.

AE is created using torch.nn.Module package and has 4 layers.Unlike ALS for AE saving and loading model is not something I am missing.Can anyone help me to find a way to save and load an auto encoder.

Solution

  • The torch module provides save and load methods.

    It saves the model as .pth files 

    torch.save(model,'model.pth')

    It also has load method

    torch.load(.pth_file)