Skip to content

Commit c2345c1

Browse files
Can-Zhaoyiheng-wang-nv
Can-Zhao
and
yiheng-wang-nv
authored
update ldm noise scheduler api for inference (#452)
Fixes # . ### Description update ldm noise scheduler api for inference ### 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: Can-Zhao <volcanofly@gmail.com> Signed-off-by: Yiheng Wang <68361391+yiheng-wang-nv@users.noreply.github.com> Signed-off-by: Yiheng Wang <vennw@nvidia.com> Co-authored-by: Yiheng Wang <68361391+yiheng-wang-nv@users.noreply.github.com> Co-authored-by: Yiheng Wang <vennw@nvidia.com>
1 parent c02d168 commit c2345c1

15 files changed

+638
-22
lines changed

‎ci/run_premerge_multi_gpu.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ verify_bundle() {
7272
bash $extra_script
7373
fi
7474
# do multi gpu based unit tests
75-
pipenv run python $(pwd)/ci/unit_tests/runner.py --b "$bundle" --dist True
75+
pipenv run torchrun $(pwd)/ci/unit_tests/runner.py --b "$bundle" --dist True
7676
remove_pipenv
7777
done
7878
else

‎ci/unit_tests/test_brats_mri_axial_slices_generative_diffusion.py‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
# Copyright (c) MONAI Consortium
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
# Unless required by applicable law or agreed to in writing, software
7+
# distributed under the License is distributed on an "AS IS" BASIS,
8+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9+
# See the License for the specific language governing permissions and
10+
# limitations under the License.
11+
12+
import json
13+
import os
14+
import shutil
15+
import sys
16+
import tempfile
17+
import unittest
18+
19+
import nibabel as nib
20+
import numpy as np
21+
from monai.bundle import ConfigWorkflow
22+
from parameterized import parameterized
23+
24+
TEST_CASE_1 = [
25+
{
26+
"bundle_root": "models/brats_mri_axial_slices_generative_diffusion",
27+
"images": "$list(sorted(glob.glob(@dataset_dir + '/image_*.nii.gz')))",
28+
"labels": "$list(sorted(glob.glob(@dataset_dir + '/label_*.nii.gz')))",
29+
"train#trainer#max_epochs": 1,
30+
"train#dataset#cache_rate": 0.0,
31+
"train_batch_size_slice": 4,
32+
}
33+
]
34+
35+
TEST_CASE_2 = [{"bundle_root": "models/brats_mri_axial_slices_generative_diffusion"}]
36+
37+
38+
def test_order(test_name1, test_name2):
39+
# specify test order.
40+
# The "train_autoencoder.json" config should be tested first in order to
41+
# produce model weights for inference, and train diffusion.
42+
def get_order(name):
43+
if "autoencoder" in name:
44+
return 1 if "train" in name else 2
45+
if "diffusion" in name:
46+
return 3 if "train" in name else 4
47+
return 5
48+
49+
return get_order(test_name1) - get_order(test_name2)
50+
51+
52+
class TestLdm2d(unittest.TestCase):
53+
def setUp(self):
54+
self.dataset_dir = tempfile.mkdtemp()
55+
dataset_size = 10
56+
input_shape = (256, 256, 112)
57+
sub_dir = os.path.join(self.dataset_dir, "Task01_BrainTumour")
58+
os.makedirs(sub_dir)
59+
data_list = []
60+
for s in range(dataset_size):
61+
test_image = np.random.randint(low=0, high=2, size=input_shape).astype(np.int8)
62+
test_label = np.random.randint(low=0, high=2, size=input_shape).astype(np.int8)
63+
image_filename = os.path.join(self.dataset_dir, f"image_{s}.nii.gz")
64+
label_filename = os.path.join(self.dataset_dir, f"label_{s}.nii.gz")
65+
nib.save(nib.Nifti1Image(test_image, np.eye(4)), image_filename)
66+
nib.save(nib.Nifti1Image(test_label, np.eye(4)), label_filename)
67+
sample_dict = {"image": image_filename, "label": label_filename}
68+
data_list.append(sample_dict)
69+
# prepare a datalist file that "monai.apps.DecathlonDataset" requires
70+
full_dict = {
71+
"name": "",
72+
"description": "",
73+
"reference": "",
74+
"licence": "",
75+
"tensorImageSize": "",
76+
"modality": "",
77+
"labels": "",
78+
"numTraining": 10,
79+
"numTest": 0,
80+
"training": data_list,
81+
}
82+
with open(os.path.join(sub_dir, "dataset.json"), "w") as f:
83+
json.dump(full_dict, f)
84+
85+
def tearDown(self):
86+
shutil.rmtree(self.dataset_dir)
87+
88+
@parameterized.expand([TEST_CASE_1])
89+
def test_autoencoder_train(self, override):
90+
override["dataset_dir"] = self.dataset_dir
91+
bundle_root = override["bundle_root"]
92+
sys.path = [bundle_root] + sys.path
93+
94+
trainer = ConfigWorkflow(
95+
workflow="train",
96+
config_file=os.path.join(bundle_root, "configs/train_autoencoder.json"),
97+
logging_file=os.path.join(bundle_root, "configs/logging.conf"),
98+
meta_file=os.path.join(bundle_root, "configs/metadata.json"),
99+
**override,
100+
)
101+
trainer.initialize()
102+
trainer.run()
103+
trainer.finalize()
104+
105+
@parameterized.expand([TEST_CASE_2])
106+
def test_autoencoder_infer(self, override):
107+
override["dataset_dir"] = self.dataset_dir
108+
bundle_root = override["bundle_root"]
109+
sys.path = [bundle_root] + sys.path
110+
111+
inferrer = ConfigWorkflow(
112+
workflow="infer",
113+
config_file=os.path.join(bundle_root, "configs/inference_autoencoder.json"),
114+
logging_file=os.path.join(bundle_root, "configs/logging.conf"),
115+
meta_file=os.path.join(bundle_root, "configs/metadata.json"),
116+
**override,
117+
)
118+
inferrer.initialize()
119+
inferrer.run()
120+
inferrer.finalize()
121+
122+
@parameterized.expand([TEST_CASE_1])
123+
def test_diffusion_train(self, override):
124+
override["dataset_dir"] = self.dataset_dir
125+
bundle_root = override["bundle_root"]
126+
sys.path = [bundle_root] + sys.path
127+
autoencoder_file = os.path.join(bundle_root, "configs/train_autoencoder.json")
128+
diffusion_file = os.path.join(bundle_root, "configs/train_diffusion.json")
129+
130+
trainer = ConfigWorkflow(
131+
workflow="train",
132+
config_file=[autoencoder_file, diffusion_file],
133+
logging_file=os.path.join(bundle_root, "configs/logging.conf"),
134+
meta_file=os.path.join(bundle_root, "configs/metadata.json"),
135+
**override,
136+
)
137+
trainer.initialize()
138+
# TODO: uncomment the following check after we have monai > 1.2.0
139+
# https://github.com/Project-MONAI/MONAI/issues/6602
140+
# check_result = trainer.check_properties()
141+
# if check_result is not None and len(check_result) > 0:
142+
# raise ValueError(f"check properties for overrided train config failed: {check_result}")
143+
trainer.run()
144+
trainer.finalize()
145+
146+
@parameterized.expand([TEST_CASE_2])
147+
def test_diffusion_infer(self, override):
148+
override["dataset_dir"] = self.dataset_dir
149+
bundle_root = override["bundle_root"]
150+
sys.path = [bundle_root] + sys.path
151+
152+
inferrer = ConfigWorkflow(
153+
workflow="infer",
154+
config_file=os.path.join(bundle_root, "configs/inference.json"),
155+
logging_file=os.path.join(bundle_root, "configs/logging.conf"),
156+
meta_file=os.path.join(bundle_root, "configs/metadata.json"),
157+
**override,
158+
)
159+
inferrer.initialize()
160+
inferrer.run()
161+
inferrer.finalize()
162+
163+
164+
if __name__ == "__main__":
165+
loader = unittest.TestLoader()
166+
loader.sortTestMethodsUsing = test_order
167+
unittest.main(testLoader=loader)

‎ci/unit_tests/test_brats_mri_axial_slices_generative_diffusion_dist.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
# Copyright (c) MONAI Consortium
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
# Unless required by applicable law or agreed to in writing, software
7+
# distributed under the License is distributed on an "AS IS" BASIS,
8+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9+
# See the License for the specific language governing permissions and
10+
# limitations under the License.
11+
12+
import json
13+
import os
14+
import shutil
15+
import tempfile
16+
import unittest
17+
18+
import nibabel as nib
19+
import numpy as np
20+
import torch
21+
from parameterized import parameterized
22+
from utils import export_config_and_run_mgpu_cmd
23+
24+
TEST_CASE_1 = [
25+
{
26+
"bundle_root": "models/brats_mri_axial_slices_generative_diffusion",
27+
"images": "$list(sorted(glob.glob(@dataset_dir + '/image_*.nii.gz')))",
28+
"labels": "$list(sorted(glob.glob(@dataset_dir + '/label_*.nii.gz')))",
29+
"train#trainer#max_epochs": 1,
30+
"train#dataset#cache_rate": 0.0,
31+
"train_batch_size_slice": 4,
32+
}
33+
]
34+
35+
36+
class TestLdm2dMGPU(unittest.TestCase):
37+
def setUp(self):
38+
self.dataset_dir = tempfile.mkdtemp()
39+
dataset_size = 10
40+
input_shape = (256, 256, 112)
41+
sub_dir = os.path.join(self.dataset_dir, "Task01_BrainTumour")
42+
os.makedirs(sub_dir)
43+
data_list = []
44+
for s in range(dataset_size):
45+
test_image = np.random.randint(low=0, high=2, size=input_shape).astype(np.int8)
46+
test_label = np.random.randint(low=0, high=2, size=input_shape).astype(np.int8)
47+
image_filename = os.path.join(self.dataset_dir, f"image_{s}.nii.gz")
48+
label_filename = os.path.join(self.dataset_dir, f"label_{s}.nii.gz")
49+
nib.save(nib.Nifti1Image(test_image, np.eye(4)), image_filename)
50+
nib.save(nib.Nifti1Image(test_label, np.eye(4)), label_filename)
51+
sample_dict = {"image": image_filename, "label": label_filename}
52+
data_list.append(sample_dict)
53+
# prepare a datalist file that "monai.apps.DecathlonDataset" requires
54+
full_dict = {
55+
"name": "",
56+
"description": "",
57+
"reference": "",
58+
"licence": "",
59+
"tensorImageSize": "",
60+
"modality": "",
61+
"labels": "",
62+
"numTraining": 10,
63+
"numTest": 0,
64+
"training": data_list,
65+
}
66+
with open(os.path.join(sub_dir, "dataset.json"), "w") as f:
67+
json.dump(full_dict, f)
68+
69+
def tearDown(self):
70+
shutil.rmtree(self.dataset_dir)
71+
72+
@parameterized.expand([TEST_CASE_1])
73+
def test_mgpu(self, override):
74+
override["dataset_dir"] = self.dataset_dir
75+
bundle_root = override["bundle_root"]
76+
autoencoder_file = os.path.join(bundle_root, "configs/train_autoencoder.json")
77+
diffusion_file = os.path.join(bundle_root, "configs/train_diffusion.json")
78+
mgpu_autoencoder_file = os.path.join(bundle_root, "configs/multi_gpu_train_autoencoder.json")
79+
mgpu_diffusion_file = os.path.join(bundle_root, "configs/multi_gpu_train_diffusion.json")
80+
n_gpu = torch.cuda.device_count()
81+
82+
export_config_and_run_mgpu_cmd(
83+
config_file=[autoencoder_file, mgpu_autoencoder_file],
84+
logging_file=os.path.join(bundle_root, "configs/logging.conf"),
85+
meta_file=os.path.join(bundle_root, "configs/metadata.json"),
86+
override_dict=override,
87+
output_path=os.path.join(bundle_root, "configs/autoencoder_override.json"),
88+
ngpu=n_gpu,
89+
)
90+
91+
export_config_and_run_mgpu_cmd(
92+
config_file=[autoencoder_file, diffusion_file, mgpu_autoencoder_file, mgpu_diffusion_file],
93+
logging_file=os.path.join(bundle_root, "configs/logging.conf"),
94+
meta_file=os.path.join(bundle_root, "configs/metadata.json"),
95+
override_dict=override,
96+
output_path=os.path.join(bundle_root, "configs/diffusion_override.json"),
97+
ngpu=n_gpu,
98+
)
99+
100+
101+
if __name__ == "__main__":
102+
unittest.main()

0 commit comments

Comments
 (0)