Skip to content

Commit

Permalink
Add an arg to specify trt output path for vista3d (#704)
Browse files Browse the repository at this point in the history
Fixes # .

### Description
A few sentences describing the changes proposed in this pull request.

### Status
**Ready/Work in progress/Hold**

### Please ensure all the checkboxes:
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Codeformat tests passed locally by running `./runtests.sh
--codeformat`.
- [ ] In-line docstrings updated.
- [ ] Update `version` and `changelog` in `metadata.json` if changing an
existing bundle.
- [ ] Please ensure the naming rules in config files meet our
requirements (please refer to: `CONTRIBUTING.md`).
- [ ] Ensure versions of packages such as `monai`, `pytorch` and `numpy`
are correct in `metadata.json`.
- [ ] Descriptions should be consistent with the content, such as
`eval_metrics` of the provided weights and TorchScript modules.
- [ ] Files larger than 25MB are excluded and replaced by providing
download links in `large_file.yml`.
- [ ] Avoid using path that contains personal information within config
files (such as use `/home/your_name/` for `"bundle_root"`).

---------

Signed-off-by: Yiheng Wang <[email protected]>
  • Loading branch information
yiheng-wang-nv authored Oct 25, 2024
1 parent dfae206 commit b445b4e
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 4 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
default_language_version:
python: python3.8
python: python3

ci:
autofix_prs: true
Expand Down
5 changes: 3 additions & 2 deletions models/vista3d/configs/inference_trt.json
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
{
"base_path": null,
"+imports": [
"$from monai.networks import trt_compile"
],
Expand All @@ -8,11 +9,11 @@
"dynamic_batchsize": "$[1, @inferer#sw_batch_size, @inferer#sw_batch_size]"
},
"network_dev": "$@network_def.to(@device)",
"encoder": "$trt_compile(@network_dev, @bundle_root + '/models/model.pt', args=@network_trt_args, submodule=['image_encoder.encoder'])",
"encoder": "$trt_compile(@network_dev, @bundle_root + '/models/model.pt' if not @base_path else @base_path, args=@network_trt_args, submodule=['image_encoder.encoder'])",
"head_trt_args": {
"dynamic_batchsize": "$[1, 1, @max_prompt_size]",
"fallback": "$True"
},
"head": "$trt_compile(@network_dev, @bundle_root + '/models/model.pt', args=@head_trt_args, submodule=['class_head']) if @head_trt_enabled else @network_dev",
"head": "$trt_compile(@network_dev, @bundle_root + '/models/model.pt' if not @base_path else @base_path, args=@head_trt_args, submodule=['class_head']) if @head_trt_enabled else @network_dev",
"network": "$None if @encoder is None else @head"
}
3 changes: 2 additions & 1 deletion models/vista3d/configs/metadata.json
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
{
"schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20240725.json",
"version": "0.5.4",
"version": "0.5.5",
"changelog": {
"0.5.5": "add arg for trt compiler base path",
"0.5.4": "add undefined label prompt check",
"0.5.3": "update readme",
"0.5.2": "fix eval issue",
Expand Down

0 comments on commit b445b4e

Please sign in to comment.