MLFlow 원격 모델 저장소에서 pytorch 모델 로딩 시 에러

 

현재 회사에서 오프라인 배치 서빙을 수행하는 스크립트를 작성하는 기회가 있었습니다. 이를 위해 타팀 연구원분께서 개발하신 딥러닝 모델을 전달받아 사용하여야 했고, 모델을 공유할 방안을 마련하여야 했습니다. 기존에는 slack이나 google drive를 사용하여 모델을 전달받았지만, 회사에 MLFlow 프로토타입을 구축하여 운용하고 있었으므로 MLFlow 원격저장소 기능을 테스트할 좋은 기회라고 생각했습니다.

 아직 딥러닝 연구원분들은 MLFlow를 잘 다루지 못하므로, 일단 slack으로 모델을 전달받고 제가 원격저장소에 업로드하였습니다. 이때 얻은 경험은 MLFlow 운영에 큰 도움이될 것 같았습니다.

 

전에 포스팅한 것 처럼, MLFlow 업로드는 그리 어렵지 않았습니다. mlflow.pytorch.log_model 메소드를 사용하면 업로드가 자동으로 되었습니다. 이때 사용한 코드는 아래와 같습니다.

from ResNeSt import ResNet

pose_model = ResNet(...)

mlflow.pytorch.log_model(
        model,
        artifact_path=experiment_name,
        registered_model_name=experiment_name,
        extra_files=extra_files
    )

 

원격 저장소에는 문제없이 업로드가 되었지만 다운로드받을 때 문제가 발생하였습니다. 분명히 튜토리얼에서는 아래와 같은 코드를 통해서 문제없이 다운로드 받았지만, 언제나 그랬듯이 실제 적용에는 여러 문제가 발생했습니다. 이 때 발생한 에러는 ResNeSt 모듈을 찾을 수 없다는 것이었습니다. 

run_id = my_run_id
model_name = my_model_name
logged_model = 'runs:/{}/{}'.format(run_id, model_name)
loaded_model = mlflow.pytorch.load_model(logged_model)

ModuleNotFoundError: No module named 'ResNeSt'

이 때 발생한 문제는 이와 같았습니다. mlflow.pytorch.load_model은 pytorch.load를 사용합니다. pytorch.load는 모델이 참조한 모듈을 확인하는데, 제가 업로드한 ResNet의 원본 소스코드는 ResNeSt.py 파일에 존재합니다. 따라서 pytorch.load를 사용해서 ResNet 모델을 불러올 때는 ResNeSt.py 소스코드를 참조하게 됩니다. 그런데 당연하게도 원격저장소에서 ResNet 모델을 불러올 때에는 ResNest파일이 존재하지 않으므로, ModuleNotFoundError와 함께 에러가 발생하게 됩니다.

 해결방안을 찾으려고 노력한 결과, 모델을 업로드할 때 ResNest.py 코드도 함께 올리고 추후 모델 로딩시에 활용하면 되겠다라고 생각하였습니다. 다행히도 log_model내의 인자에서 extra_files라는 키워드로 추가로 업로드할 파일을 선택할 수 있는 기능이 있었습니다. 따라서 아래의 코드와 같이 log_model 함수에 extra_files 인자를 추가하였습니다.

mlflow.pytorch.log_model(
        model,
        artifact_path=experiment_name,
        registered_model_name=experiment_name,
        extra_files=['ResNeSt.py']
    )

그리고 다시 아래의 코드를 사용하여 모델을 로딩하였는데.. 똑같이 ModuleNotFoundError가 발생합니다. 모델 로딩시에 자동으로 extra_file를 활용하지 않는 것 같았습니다.

run_id = my_run_id
model_name = my_model_name
logged_model = 'runs:/{}/{}'.format(run_id, model_name)
loaded_model = mlflow.pytorch.load_model(logged_model)

ModuleNotFoundError: No module named 'ResNeSt'

 

따라서 수동으로 모델 아티팩트를 저장한 후 sys.path.append로 extra_file이 위치한 디렉토리를 추가하는 방법으로 문제를 해결했습니다.

run_id = my_run_id
model_file_path=PROJECT_PATH+'/model/model_archive'
 mlflow_client.download_artifacts(run_id, model_name, model_file_path)
 model_extra_files_path = model_file_path + '/extra_files'.format(model_name)

 if os.path.isdir(model_extra_files_path):
     sys.path.append(model_extra_files_path)

 logged_model = 'runs:/{}/{}'.format(run_id, model_name)
 loaded_model = mlflow.pytorch.load_model(logged_model)

이렇게 코드를 작성하니 더이상 ModuleNotFoundError가 발생하지 않았습니다..

 

마지막으로...

아직 MLOps가 널리 퍼지지 않은 상태이고, 따라서 MLFlow의 사용자도 많지 않은 것 같습니다. 아마 다른 MLOps 분들도 저와 비슷한 문제 혹은 고민이 있으실 것 같습니다. 앞으로 제가 직면한 문제에 대해 계속 포스팅을 할 예정이니, 혹시 필요하신 분들은 참고하셔도 좋을 것 같습니다.