diff --git a/docker/serve/Dockerfile b/docker/serve/Dockerfile index c50c4e8e..3ea8849b 100644 --- a/docker/serve/Dockerfile +++ b/docker/serve/Dockerfile @@ -1,6 +1,7 @@ ARG PYTORCH="2.0.1" ARG CUDA="11.7" ARG CUDNN="8" +ARG TORCHSERVE="0.7.1" FROM pytorch/torchserve:latest-gpu ARG MMPRE="1.2.0" @@ -11,8 +12,8 @@ ENV HOME="/home/model-server" ENV PATH="/opt/conda/bin:$HOME/.local/bin:$PATH" RUN export FORCE_CUDA=1 -# TORCHSEVER -RUN pip install torchserve torch-model-archiver +# TORCHSERVE +RUN pip install torchserve==${TORCHSERVE} torch-model-archiver==${TORCHSERVE} RUN pip install nvgpu # OPEN-MMLAB diff --git a/tools/torchserve/mmpretrain2torchserve.py b/tools/torchserve/mmpretrain2torchserve.py index 8d53bf3f..a0e0384a 100644 --- a/tools/torchserve/mmpretrain2torchserve.py +++ b/tools/torchserve/mmpretrain2torchserve.py @@ -10,7 +10,8 @@ try: from model_archiver.model_packaging_utils import ModelExportUtils except ImportError: raise ImportError( - 'Please run `pip install torchserve torch-model-archiver"` to ' + 'Please run ' + '`pip install torchserve==0.7.1 torch-model-archiver==0.7.1"` to ' 'install required third-party libraries.') @@ -106,7 +107,7 @@ if __name__ == '__main__': if package_model is None: raise ImportError('`torch-model-archiver` is required.' - 'Try: pip install torch-model-archiver') + 'Try: pip install torch-model-archiver==0.7.1') mmpretrain2torchserve(args.config, args.checkpoint, args.output_folder, args.model_name, args.model_version, args.force)