Hello, I have used mlflow
for logging my experiments, aswell as my models. The model that I am using is pytorch-forecasting
’ TemporalFusionTransformer
model trained on kaggle/colab gpus. After my experimentation, I needed to load the model from S3 bucket where I was storing my runs and use my selected model for predictions. However, when I try to load the model when writing the production script I get the error: AttributeError: 'PassThroughProfiler' object has no attribute '_stage'
.
My producion script (so far) is this:
import os
import mlflow
MLFLOW_TRACKING_URI='model-uri'
AWS_ACCESS_KEY_ID = 'aws-access-key-id'
AWS_SECRET_ACCESS_KEY = 'aws-secret-acces-key'
MODEL_URI = 's3:path'
os.environ["AWS_ACCESS_KEY_ID"] = AWS_ACCESS_KEY_ID
os.environ["AWS_SECRET_ACCESS_KEY"] = AWS_SECRET_ACCESS_KEY
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
print('tmp')
if __name__ == '__main__':
model = mlflow.pytorch.load_model(MODEL_URI)
The full stack trace is this:
Traceback (most recent call last):
File "/home/petar/Praksa/Production script/production_script.py", line 24, in <module>
model = mlflow.pytorch.load_model(MODEL_URI)
File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/mlflow/pytorch/__init__.py", line 714, in load_model
return _load_model(path=torch_model_artifacts_path, **kwargs)
File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/mlflow/pytorch/__init__.py", line 626, in _load_model
return torch.load(model_path, **kwargs)
File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
result = unpickler.load()
TypeError: an integer is required (got type bytes)
Exception ignored in: <function Profiler.__del__ at 0x7f98d8c59040>
Traceback (most recent call last):
File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/pytorch_lightning/profilers/profiler.py", line 170, in __del__
self.teardown(stage=self._stage)
AttributeError: 'PassThroughProfiler' object has no attribute '_stage'
In the stack trace above, it says that the error is thrown in the profiler of pytorch-lightning
.
However, if I try to load the same model in a kaggle notebook it runs fine and returns the model.
The dependencies used in the notebook and in my script are the same and are as follows:
-f https://download.pytorch.org/whl/torch_stable.html
pytorch-forecasting==0.9.0
protobuf==3.20.0
torch==1.9.0+cu111
torchvision==0.10.0+cu111
torchtext==0.10.0
pytorch-lightning==1.7.7
pandas==1.2.5
numpy==1.23.4
mlflow==1.21.0
scikit-learn==0.24.2
psycopg2-binary==2.9.2
boto3
Any help towards resolving this issue is welcomed.