Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Benchmark VISTA3D #671

Draft
wants to merge 2 commits into
base: dev
Choose a base branch
from
Draft

Benchmark VISTA3D #671

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f1a3af4
add range function to the bundle
Sep 23, 2024
501911c
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 57 additions & 38 deletions models/vista3d/configs/inference.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@
"output_dtype": "$np.float32",
"output_postfix": "trans",
"separate_folder": true,
"input_dict": "${'image': '/data/Task09_Spleen/imagesTr/spleen_10.nii.gz', 'label_prompt': [3]}",
"sw_batch_size": 10,
Copy link
Collaborator

@SachidanandAlle SachidanandAlle Sep 23, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we determine it based on GPU memory size?
torch.cuda.get_device_properties('device').total_memory will give you the total memory. based on total memory and patch size you can come with a simple logic to determine if you want to use batch size 1,2,4,8,10?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @SachidanandAlle , I am not sure what the a simple logic to determine if you want to use batch size 1,2,4,8,10 means here. If the patch size is fixed, then GPU memory usage is fixed. Do you suggest to change to a smaller batch size if the GPU memory is not enough? Otherwise I think we should always choose the 14 as the batchsize here.

Thanks
Bin

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

based on GPU size determine what is the best batch size you can use. if you have enough gpu memory then use higher batch sizes. have you tried running batch size 10 for 16 GB v100? if batch size 10 works for a typical 512x512x256 image then, we are good with default 10 as the batch size

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also it shall be great if we can modify our sw_inferer logic to determine (we can be little smart here) the required sw_batch size based on the following 3 things.

  1. Current Input Image Size (after all pre-processing etc...)
  2. GPU Memory size (max capacity)
  3. Patch Size (typically recommended by the trained model)

cc: @Nic-Ma

Copy link
Contributor Author

@binliunls binliunls Sep 25, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @SachidanandAlle ,

I've updated the benchmark results with more images. According to some results, the latency will reach the best one at a specific batch size. Both increase and decrease will make the latency worse. I think this batch size is a tradeoff betweeen overhead and calculation. When batchsize is small, the overhead will waste too much time. But when batchsize is large, the calculation will waste too much time.

And I also checked the GPU memory this time. I note that for some images even batchsize 1 exceed the 16GB for V100 16GB GPU. You can check the detail in the image shown below. I don't think VISTA3D is a model that suit for GPU memory size less than 16GB. And for acceleration purpose, we may need more GPU memory.

image

I am trying to benchmark a (512, 512, 256) image. Will update the result later.

"input_dict": "${'image': '/home/liubin/data/trt_2408/dataset/Task09_Spleen/imagesTr/spleen_10.nii.gz'}",
"everything_labels": "$list(set([i+1 for i in range(132)]) - set([2,16,18,20,21,23,24,25,26,27,128,129,130,131,132]))",
"metadata_path": "$@bundle_root + '/configs/metadata.json'",
"metadata": "$json.loads(pathlib.Path(@metadata_path).read_text())",
Expand All @@ -40,7 +41,6 @@
1.5,
1.5
],
"sw_batch_size": 1,
"patch_size": [
128,
128,
Expand Down Expand Up @@ -105,9 +105,19 @@
"dtype": "$torch.float32"
}
],
"range_preprocessing": {
"_target_": "Range",
"name": "preprocessing",
"recursive": true
},
"range_postprocessing": {
"_target_": "Range",
"name": "postprocessing",
"recursive": true
},
"preprocessing": {
"_target_": "Compose",
"transforms": "$@preprocessing_transforms "
"transforms": "$@range_preprocessing(@preprocessing_transforms)"
},
"dataset": {
"_target_": "Dataset",
Expand All @@ -128,48 +138,57 @@
"sw_batch_size": "@sw_batch_size",
"use_point_window": "@use_point_window"
},
"postprocessing_transforms": [
{
"_target_": "ToDeviced",
"keys": "pred",
"device": "cpu",
"_disabled_": true
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since this is disabled - why don't we remove this transform from post?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@heyufan1995 could you please answer this one?

Thanks
Bin

},
{
"_target_": "monai.apps.vista3d.transforms.VistaPostTransformd",
"keys": "pred"
},
{
"_target_": "Invertd",
"keys": "pred",
"transform": "$copy.deepcopy(@preprocessing)",
"orig_keys": "@image_key",
"nearest_interp": true,
"to_tensor": true
},
{
"_target_": "Lambdad",
"func": "$lambda x: torch.nan_to_num(x, nan=255)",
"keys": "pred"
},
{
"_target_": "SaveImaged",
"keys": "pred",
"resample": false,
"output_dir": "@output_dir",
"output_ext": "@output_ext",
"output_dtype": "@output_dtype",
"output_postfix": "@output_postfix",
"separate_folder": "@separate_folder"
}
],
"postprocessing": {
"_target_": "Compose",
"transforms": [
{
"_target_": "ToDeviced",
"keys": "pred",
"device": "cpu",
"_disabled_": true
},
{
"_target_": "monai.apps.vista3d.transforms.VistaPostTransformd",
"keys": "pred"
},
{
"_target_": "Invertd",
"keys": "pred",
"transform": "$copy.deepcopy(@preprocessing)",
"orig_keys": "@image_key",
"nearest_interp": true,
"to_tensor": true
},
{
"_target_": "Lambdad",
"func": "$lambda x: torch.nan_to_num(x, nan=255)",
"keys": "pred"
},
{
"_target_": "SaveImaged",
"keys": "pred",
"resample": false,
"output_dir": "@output_dir",
"output_ext": "@output_ext",
"output_dtype": "@output_dtype",
"output_postfix": "@output_postfix",
"separate_folder": "@separate_folder"
}
]
"transforms": "$@range_postprocessing(@postprocessing_transforms)"
},
"handlers": [
{
"_target_": "StatsHandler",
"iteration_log": false
},
{
"_target_": "RangeHandler",
"events": "ITERATION"
},
{
"_target_": "RangeHandler",
"events": "BATCH"
}
],
"checkpointloader": {
Expand Down
69 changes: 36 additions & 33 deletions models/vista3d/scripts/inferer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import torch
from monai.apps.vista3d.inferer import point_based_window_inferer
from monai.inferers import Inferer, SlidingWindowInfererAdapt
from monai.utils import Range
from torch import Tensor


Expand Down Expand Up @@ -80,38 +81,40 @@ def __call__(
device = inputs[0].device
else:
device = inputs.device
val_outputs = point_based_window_inferer(
inputs=inputs,
roi_size=self.roi_size,
sw_batch_size=self.sw_batch_size,
transpose=True,
with_coord=True,
predictor=network,
mode="gaussian",
sw_device=device,
device=device,
overlap=self.overlap,
point_coords=point_coords,
point_labels=point_labels,
class_vector=class_vector,
prompt_class=prompt_class,
prev_mask=prev_mask,
labels=labels,
label_set=label_set,
)
with Range("SW_PointInfer"):
val_outputs = point_based_window_inferer(
inputs=inputs,
roi_size=self.roi_size,
sw_batch_size=self.sw_batch_size,
transpose=True,
with_coord=True,
predictor=network,
mode="gaussian",
sw_device=device,
device=device,
overlap=self.overlap,
point_coords=point_coords,
point_labels=point_labels,
class_vector=class_vector,
prompt_class=prompt_class,
prev_mask=prev_mask,
labels=labels,
label_set=label_set,
)
else:
val_outputs = SlidingWindowInfererAdapt(
roi_size=self.roi_size, sw_batch_size=self.sw_batch_size, with_coord=True
)(
inputs,
network,
transpose=True,
point_coords=point_coords,
point_labels=point_labels,
class_vector=class_vector,
prompt_class=prompt_class,
prev_mask=prev_mask,
labels=labels,
label_set=label_set,
)
with Range("SW_Infer"):
val_outputs = SlidingWindowInfererAdapt(
roi_size=self.roi_size, sw_batch_size=self.sw_batch_size, with_coord=True
)(
inputs,
network,
transpose=True,
point_coords=point_coords,
point_labels=point_labels,
class_vector=class_vector,
prompt_class=prompt_class,
prev_mask=prev_mask,
labels=labels,
label_set=label_set,
)
return val_outputs
Loading