-
Notifications
You must be signed in to change notification settings - Fork 1.6k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(components): PyTorch - Added the Create PyTorch Model Archive co…
…mponent (#5630) * Components - PyTorch - Added the Create PyTorch Model Archive component * Added component to the sample pipeline
- Loading branch information
Showing
2 changed files
with
51 additions
and
0 deletions.
There are no files selected for viewing
39 changes: 39 additions & 0 deletions
39
components/PyTorch/Create_PyTorch_Model_Archive/component.yaml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,39 @@ | ||
name: Create PyTorch Model Archive | ||
inputs: | ||
- {name: Model, type: PyTorchScriptModule} | ||
- {name: Model name, type: String, default: model} | ||
- {name: Model version, type: String, default: "1.0"} | ||
- {name: Handler, type: PythonCode, description: "See https://github.com/pytorch/serve/blob/master/docs/custom_service.md"} | ||
outputs: | ||
- {name: Model archive, type: PyTorchModelArchive} | ||
metadata: | ||
annotations: | ||
author: Alexey Volkov <alexey.volkov@ark-kun.com> | ||
implementation: | ||
container: | ||
image: pytorch/torchserve:0.3.0-cpu | ||
command: | ||
- bash | ||
- -exc | ||
- | | ||
model_path=$0 | ||
handler_path=$1 | ||
model_name=$2 | ||
model_version=$3 | ||
output_model_archive_path=$4 | ||
mkdir -p "$(dirname "$output_model_archive_path")" | ||
# torch-model-archiver needs the handler to have .py extension | ||
cp "$handler_path" handler.py | ||
torch-model-archiver --model-name "$model_name" --version "$model_version" --serialized-file "$model_path" --handler handler.py | ||
# torch-model-archiver does not allow specifying the output path, but always writes to "${model_name}.<format>" | ||
expected_model_archive_path="${model_name}.mar" | ||
mv "$expected_model_archive_path" "$output_model_archive_path" | ||
- {inputPath: Model} | ||
- {inputPath: Handler} | ||
- {inputValue: Model name} | ||
- {inputValue: Model version} | ||
- {outputPath: Model archive} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters