monai-weekly 1.5.dev2506__py3-none-any.whl → 1.5.dev2508__py3-none-any.whl

tests/integration/test_integration_segmentation_3d.py

@@ -0,0 +1,304 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import shutil
+import tempfile
+import unittest
+from glob import glob
+
+import nibabel as nib
+import numpy as np
+import torch
+
+import monai
+from monai.data import create_test_image_3d, decollate_batch
+from monai.inferers import sliding_window_inference
+from monai.metrics import DiceMetric
+from monai.networks import eval_mode
+from monai.networks.nets import UNet
+from monai.transforms import (
+    Activations,
+    AsDiscrete,
+    Compose,
+    EnsureChannelFirstd,
+    LoadImaged,
+    RandCropByPosNegLabeld,
+    RandRotate90d,
+    SaveImage,
+    ScaleIntensityd,
+    Spacingd,
+)
+from monai.utils import optional_import, set_determinism
+from monai.visualize import plot_2d_or_3d_image
+from tests.test_utils import DistTestCase, TimedCall, skip_if_quick
+from tests.testing_data.integration_answers import test_integration_value
+
+SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter")
+
+TASK = "integration_segmentation_3d"
+
+
+def run_training_test(root_dir, device="cuda:0", cachedataset=0, readers=(None, None)):
+    monai.config.print_config()
+    images = sorted(glob(os.path.join(root_dir, "img*.nii.gz")))
+    segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
+    train_files = [{"img": img, "seg": seg} for img, seg in zip(images[:20], segs[:20])]
+    val_files = [{"img": img, "seg": seg} for img, seg in zip(images[-20:], segs[-20:])]
+
+    # define transforms for image and segmentation
+    train_transforms = Compose(
+        [
+            LoadImaged(keys=["img", "seg"], reader=readers[0]),
+            EnsureChannelFirstd(keys=["img", "seg"]),
+            # resampling with align_corners=True or dtype=float64 will generate
+            # slight different results between PyTorch 1.5 an 1.6
+            Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
+            ScaleIntensityd(keys="img"),
+            RandCropByPosNegLabeld(
+                keys=["img", "seg"], label_key="seg", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
+            ),
+            RandRotate90d(keys=["img", "seg"], prob=0.8, spatial_axes=[0, 2]),
+        ]
+    )
+    train_transforms.set_random_state(1234)
+    val_transforms = Compose(
+        [
+            LoadImaged(keys=["img", "seg"], reader=readers[1]),
+            EnsureChannelFirstd(keys=["img", "seg"]),
+            # resampling with align_corners=True or dtype=float64 will generate
+            # slight different results between PyTorch 1.5 an 1.6
+            Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
+            ScaleIntensityd(keys="img"),
+        ]
+    )
+
+    # create a training data loader
+    if cachedataset == 2:
+        train_ds = monai.data.CacheDataset(
+            data=train_files, transform=train_transforms, cache_rate=0.8, runtime_cache="process"
+        )
+    elif cachedataset == 3:
+        train_ds = monai.data.LMDBDataset(data=train_files, transform=train_transforms, cache_dir=root_dir)
+    else:
+        train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
+    # use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
+    train_loader = monai.data.DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)
+    # create a validation data loader
+    val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
+    val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
+    val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+    dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
+
+    # create UNet, DiceLoss and Adam optimizer
+    model = monai.networks.nets.UNet(
+        spatial_dims=3,
+        in_channels=1,
+        out_channels=1,
+        channels=(16, 32, 64, 128, 256),
+        strides=(2, 2, 2, 2),
+        num_res_units=2,
+    ).to(device)
+    loss_function = monai.losses.DiceLoss(sigmoid=True)
+    optimizer = torch.optim.Adam(model.parameters(), 5e-4)
+
+    # start a typical PyTorch training
+    val_interval = 2
+    best_metric, best_metric_epoch = -1, -1
+    epoch_loss_values = []
+    metric_values = []
+    writer = SummaryWriter(log_dir=os.path.join(root_dir, "runs"))
+    model_filename = os.path.join(root_dir, "best_metric_model.pth")
+    for epoch in range(6):
+        print("-" * 10)
+        print(f"Epoch {epoch + 1}/{6}")
+        model.train()
+        epoch_loss = 0
+        step = 0
+        for batch_data in train_loader:
+            step += 1
+            inputs, labels = batch_data["img"].to(device), batch_data["seg"].to(device)
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = loss_function(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item()
+            epoch_len = len(train_ds) // train_loader.batch_size
+            print(f"{step}/{epoch_len}, train_loss:{loss.item():0.4f}")
+            writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
+        epoch_loss /= step
+        epoch_loss_values.append(epoch_loss)
+        print(f"epoch {epoch + 1} average loss:{epoch_loss:0.4f}")
+
+        if (epoch + 1) % val_interval == 0:
+            with eval_mode(model):
+                val_images = None
+                val_labels = None
+                val_outputs = None
+                for val_data in val_loader:
+                    val_images, val_labels = (val_data["img"].to(device), val_data["seg"].to(device))
+                    sw_batch_size, roi_size = 4, (96, 96, 96)
+                    val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
+                    # decollate prediction into a list and execute post processing for every item
+                    val_outputs = [val_post_tran(i) for i in decollate_batch(val_outputs)]
+                    # compute metrics
+                    dice_metric(y_pred=val_outputs, y=val_labels)
+
+                metric = dice_metric.aggregate().item()
+                dice_metric.reset()
+                metric_values.append(metric)
+                if metric > best_metric:
+                    best_metric = metric
+                    best_metric_epoch = epoch + 1
+                    torch.save(model.state_dict(), model_filename)
+                    print("saved new best metric model")
+                print(
+                    f"current epoch {epoch + 1} current mean dice: {metric:0.4f} "
+                    f"best mean dice: {best_metric:0.4f} at epoch {best_metric_epoch}"
+                )
+                writer.add_scalar("val_mean_dice", metric, epoch + 1)
+                # plot the last model output as GIF image in TensorBoard with the corresponding image and label
+                plot_2d_or_3d_image(val_images, epoch + 1, writer, index=0, tag="image")
+                plot_2d_or_3d_image(val_labels, epoch + 1, writer, index=0, tag="label")
+                plot_2d_or_3d_image(val_outputs, epoch + 1, writer, index=0, tag="output")
+    print(f"train completed, best_metric: {best_metric:0.4f}  at epoch: {best_metric_epoch}")
+    writer.close()
+    return epoch_loss_values, best_metric
+
+
+def run_inference_test(root_dir, device="cuda:0"):
+    images = sorted(glob(os.path.join(root_dir, "im*.nii.gz")))
+    segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
+    val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
+
+    saver = SaveImage(
+        output_dir=os.path.join(root_dir, "output"),
+        dtype=np.float32,
+        output_ext=".nii.gz",
+        output_postfix="seg",
+        mode="bilinear",
+    )
+    # define transforms for image and segmentation
+    val_transforms = Compose(
+        [
+            LoadImaged(keys=["img", "seg"]),
+            EnsureChannelFirstd(keys=["img", "seg"]),
+            # resampling with align_corners=True or dtype=float64 will generate
+            # slight different results between PyTorch 1.5 an 1.6
+            Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
+            ScaleIntensityd(keys="img"),
+        ]
+    )
+    val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
+    # sliding window inference need to input 1 image in every iteration
+    val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
+    val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5), saver])
+    dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
+
+    model = UNet(
+        spatial_dims=3,
+        in_channels=1,
+        out_channels=1,
+        channels=(16, 32, 64, 128, 256),
+        strides=(2, 2, 2, 2),
+        num_res_units=2,
+    ).to(device)
+
+    model_filename = os.path.join(root_dir, "best_metric_model.pth")
+    model.load_state_dict(torch.load(model_filename))
+    with eval_mode(model):
+        # resampling with align_corners=True or dtype=float64 will generate
+        # slight different results between PyTorch 1.5 an 1.6
+        for val_data in val_loader:
+            val_images, val_labels = (val_data["img"].to(device), val_data["seg"].to(device))
+            # define sliding window size and batch size for windows inference
+            sw_batch_size, roi_size = 4, (96, 96, 96)
+            val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
+            # decollate prediction into a list
+            val_outputs = [val_post_tran(i) for i in decollate_batch(val_outputs)]
+            # compute metrics
+            dice_metric(y_pred=val_outputs, y=val_labels)
+
+    return dice_metric.aggregate().item()
+
+
+@skip_if_quick
+class IntegrationSegmentation3D(DistTestCase):
+    def setUp(self):
+        set_determinism(seed=0)
+
+        self.data_dir = tempfile.mkdtemp()
+        for i in range(40):
+            im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1, channel_dim=-1)
+            n = nib.Nifti1Image(im, np.eye(4))
+            nib.save(n, os.path.join(self.data_dir, f"img{i:d}.nii.gz"))
+            n = nib.Nifti1Image(seg, np.eye(4))
+            nib.save(n, os.path.join(self.data_dir, f"seg{i:d}.nii.gz"))
+
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu:0"
+
+    def tearDown(self):
+        set_determinism(seed=None)
+        shutil.rmtree(self.data_dir)
+
+    def train_and_infer(self, idx=0):
+        results = []
+        set_determinism(0)
+        _readers = (None, None)
+        if idx == 1:
+            _readers = ("itkreader", "itkreader")
+        elif idx == 2:
+            _readers = ("itkreader", "nibabelreader")
+        losses, best_metric = run_training_test(self.data_dir, device=self.device, cachedataset=idx, readers=_readers)
+        infer_metric = run_inference_test(self.data_dir, device=self.device)
+
+        # check training properties
+        print("losses", losses)
+        print("best metric", best_metric)
+        print("infer metric", infer_metric)
+        self.assertTrue(len(glob(os.path.join(self.data_dir, "runs"))) > 0)
+        model_file = os.path.join(self.data_dir, "best_metric_model.pth")
+        self.assertTrue(os.path.exists(model_file))
+
+        # check inference properties
+        output_files = sorted(glob(os.path.join(self.data_dir, "output", "img*", "*.nii.gz")))
+        print([np.mean(nib.load(output).get_fdata()) for output in output_files])
+        results.extend(losses)
+        results.append(best_metric)
+        results.append(infer_metric)
+        for output in output_files:
+            ave = np.mean(nib.load(output).get_fdata())
+            results.append(ave)
+        self.assertTrue(test_integration_value(TASK, key="losses", data=results[:6], rtol=1e-3))
+        self.assertTrue(test_integration_value(TASK, key="best_metric", data=results[6], rtol=1e-2))
+        self.assertTrue(test_integration_value(TASK, key="infer_metric", data=results[7], rtol=1e-2))
+        self.assertTrue(test_integration_value(TASK, key="output_sums", data=results[8:], rtol=5e-2))
+        return results
+
+    def test_training(self):
+        repeated = []
+        for i in range(4):
+            results = self.train_and_infer(i)
+            repeated.append(results)
+        np.testing.assert_allclose(repeated[0], repeated[1])
+        np.testing.assert_allclose(repeated[0], repeated[2])
+        np.testing.assert_allclose(repeated[0], repeated[3])
+
+    @TimedCall(seconds=360, daemon=False)
+    def test_timing(self):
+        self.train_and_infer(idx=3)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/integration/test_integration_sliding_window.py

@@ -0,0 +1,100 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import tempfile
+import unittest
+
+import nibabel as nib
+import numpy as np
+import torch
+from ignite.engine import Engine, Events
+from torch.utils.data import DataLoader
+
+from monai.data import ImageDataset, create_test_image_3d
+from monai.inferers import sliding_window_inference
+from monai.networks import eval_mode, predict_segmentation
+from monai.networks.nets import UNet
+from monai.transforms import EnsureChannelFirst, SaveImage
+from monai.utils import set_determinism
+from tests.test_utils import DistTestCase, TimedCall, make_nifti_image, skip_if_quick
+
+
+def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"):
+    ds = ImageDataset(
+        [img_name],
+        [seg_name],
+        transform=EnsureChannelFirst(channel_dim="no_channel"),
+        seg_transform=EnsureChannelFirst(channel_dim="no_channel"),
+        image_only=True,
+    )
+    loader = DataLoader(ds, batch_size=1, pin_memory=torch.cuda.is_available())
+
+    net = UNet(
+        spatial_dims=3, in_channels=1, out_channels=1, channels=(4, 8, 16, 32), strides=(2, 2, 2), num_res_units=2
+    ).to(device)
+    roi_size = (16, 32, 48)
+    sw_batch_size = batch_size
+
+    saver = SaveImage(output_dir=output_dir, output_ext=".nii.gz", output_postfix="seg")
+
+    def _sliding_window_processor(_engine, batch):
+        img = batch[0]  # first item from ImageDataset is the input image
+        with eval_mode(net):
+            seg_probs = sliding_window_inference(img.to(device), roi_size, sw_batch_size, net, device=device)
+            return predict_segmentation(seg_probs)
+
+    def save_func(engine):
+        for m in engine.state.output:
+            saver(m)
+
+    infer_engine = Engine(_sliding_window_processor)
+    infer_engine.add_event_handler(Events.ITERATION_COMPLETED, save_func)
+    infer_engine.run(loader)
+
+    basename = os.path.basename(img_name)[: -len(".nii.gz")]
+    saved_name = os.path.join(output_dir, basename, f"{basename}_seg.nii.gz")
+    return saved_name
+
+
+@skip_if_quick
+class TestIntegrationSlidingWindow(DistTestCase):
+    def setUp(self):
+        set_determinism(seed=0)
+
+        im, seg = create_test_image_3d(28, 25, 63, rad_max=10, noise_max=1, num_objs=4, num_seg_classes=1)
+        self.img_name = make_nifti_image(im)
+        self.seg_name = make_nifti_image(seg)
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
+
+    def tearDown(self):
+        set_determinism(seed=None)
+        if os.path.exists(self.img_name):
+            os.remove(self.img_name)
+        if os.path.exists(self.seg_name):
+            os.remove(self.seg_name)
+
+    @TimedCall(seconds=20)
+    def test_training(self):
+        set_determinism(seed=0)
+        with tempfile.TemporaryDirectory() as tempdir:
+            output_file = run_test(
+                batch_size=2, img_name=self.img_name, seg_name=self.seg_name, output_dir=tempdir, device=self.device
+            )
+            output_image = nib.load(output_file).get_fdata()
+            np.testing.assert_allclose(np.sum(output_image), 33621)
+            np.testing.assert_allclose(output_image.shape, (28, 25, 63))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/integration/test_integration_stn.py

@@ -0,0 +1,133 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+from monai.data import create_test_image_2d
+from monai.networks.layers import AffineTransform
+from monai.utils import set_determinism
+from tests.test_utils import DistTestCase, TimedCall
+
+
+class STNBenchmark(nn.Module):
+    """
+    adapted from https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html
+    """
+
+    def __init__(self, is_ref=True, reverse_indexing=False):
+        super().__init__()
+        self.is_ref = is_ref
+        self.localization = nn.Sequential(
+            nn.Conv2d(1, 8, kernel_size=7),
+            nn.MaxPool2d(2, stride=2),
+            nn.ReLU(True),
+            nn.Conv2d(8, 10, kernel_size=5),
+            nn.MaxPool2d(2, stride=2),
+            nn.ReLU(True),
+        )
+        # Regressor for the 3 * 2 affine matrix
+        self.fc_loc = nn.Sequential(nn.Linear(10 * 3 * 3, 32), nn.ReLU(True), nn.Linear(32, 3 * 2))
+        # Initialize the weights/bias with identity transformation
+        self.fc_loc[2].weight.data.zero_()
+        self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
+        if not self.is_ref:
+            self.xform = AffineTransform(align_corners=False, normalized=True, reverse_indexing=reverse_indexing)
+
+    # Spatial transformer network forward function
+    def stn_ref(self, x):
+        xs = self.localization(x)
+        xs = xs.view(-1, 10 * 3 * 3)
+        theta = self.fc_loc(xs)
+        theta = theta.view(-1, 2, 3)
+
+        grid = F.affine_grid(theta, x.size(), align_corners=False)
+        x = F.grid_sample(x, grid, align_corners=False)
+        return x
+
+    def stn(self, x):
+        xs = self.localization(x)
+        xs = xs.view(-1, 10 * 3 * 3)
+        theta = self.fc_loc(xs)
+        theta = theta.view(-1, 2, 3)
+        x = self.xform(x, theta, spatial_size=x.size()[2:])
+        return x
+
+    def forward(self, x):
+        if self.is_ref:
+            return self.stn_ref(x)
+        return self.stn(x)
+
+
+def compare_2d(is_ref=True, device=None, reverse_indexing=False):
+    batch_size = 32
+    img_a = [create_test_image_2d(28, 28, 5, rad_max=6, noise_max=1)[0][None] for _ in range(batch_size)]
+    img_b = [create_test_image_2d(28, 28, 5, rad_max=6, noise_max=1)[0][None] for _ in range(batch_size)]
+    img_a = np.stack(img_a, axis=0)
+    img_b = np.stack(img_b, axis=0)
+    img_a = torch.as_tensor(img_a, device=device)
+    img_b = torch.as_tensor(img_b, device=device)
+    model = STNBenchmark(is_ref=is_ref, reverse_indexing=reverse_indexing).to(device)
+    optimizer = optim.SGD(model.parameters(), lr=0.001)
+    model.train()
+    init_loss = None
+    for _ in range(20):
+        optimizer.zero_grad()
+        output_a = model(img_a)
+        loss = torch.mean((output_a - img_b) ** 2)
+        if init_loss is None:
+            init_loss = loss.item()
+        loss.backward()
+        optimizer.step()
+    return model(img_a).detach().cpu().numpy(), loss.item(), init_loss
+
+
+class TestSpatialTransformerCore(DistTestCase):
+    def setUp(self):
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
+
+    def tearDown(self):
+        set_determinism(seed=None)
+
+    @TimedCall(seconds=100, skip_timing=not torch.cuda.is_available())
+    def test_training(self):
+        """
+        check that the quality AffineTransform backpropagation
+        """
+        atol = 1e-5
+        set_determinism(seed=0)
+        out_ref, loss_ref, init_loss_ref = compare_2d(True, self.device)
+        print(out_ref.shape, loss_ref, init_loss_ref)
+
+        set_determinism(seed=0)
+        out, loss, init_loss = compare_2d(False, self.device)
+        print(out.shape, loss, init_loss)
+        np.testing.assert_allclose(out_ref, out, atol=atol)
+        np.testing.assert_allclose(init_loss_ref, init_loss, atol=atol)
+        np.testing.assert_allclose(loss_ref, loss, atol=atol)
+
+        set_determinism(seed=0)
+        out, loss, init_loss = compare_2d(False, self.device, True)
+        print(out.shape, loss, init_loss)
+        np.testing.assert_allclose(out_ref, out, atol=atol)
+        np.testing.assert_allclose(init_loss_ref, init_loss, atol=atol)
+        np.testing.assert_allclose(loss_ref, loss, atol=atol)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/integration/test_integration_unet_2d.py

@@ -0,0 +1,67 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import torch
+from ignite.engine import create_supervised_trainer
+from torch.utils.data import DataLoader, Dataset
+
+from monai.data import create_test_image_2d
+from monai.losses import DiceLoss
+from monai.networks.nets import BasicUNet, UNet
+from tests.test_utils import DistTestCase, TimedCall, skip_if_quick
+
+
+def run_test(net_name="basicunet", batch_size=64, train_steps=100, device="cuda:0"):
+    class _TestBatch(Dataset):
+        def __getitem__(self, _unused_id):
+            im, seg = create_test_image_2d(128, 128, noise_max=1, num_objs=4, num_seg_classes=1)
+            return im[None], seg[None].astype(np.float32)
+
+        def __len__(self):
+            return train_steps
+
+    net = None
+    if net_name == "basicunet":
+        net = BasicUNet(spatial_dims=2, in_channels=1, out_channels=1, features=(4, 8, 8, 16, 16, 32))
+    elif net_name == "unet":
+        net = UNet(
+            spatial_dims=2, in_channels=1, out_channels=1, channels=(4, 8, 16, 32), strides=(2, 2, 2), num_res_units=2
+        )
+    net.to(device)
+
+    loss = DiceLoss(sigmoid=True)
+    opt = torch.optim.Adam(net.parameters(), 1e-4)
+    src = DataLoader(_TestBatch(), batch_size=batch_size)
+
+    trainer = create_supervised_trainer(net, opt, loss, device, False)
+
+    trainer.run(src, 1)
+    loss = trainer.state.output
+    return loss
+
+
+@skip_if_quick
+class TestIntegrationUnet2D(DistTestCase):
+    @TimedCall(seconds=20, daemon=False)
+    def test_unet_training(self):
+        for n in ["basicunet", "unet"]:
+            loss = run_test(net_name=n, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0"))
+            print(loss)
+            self.assertGreaterEqual(0.85, loss)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/integration/test_integration_workers.py

@@ -0,0 +1,61 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import torch
+
+from monai.data import DataLoader
+from monai.utils import set_determinism
+from tests.test_utils import DistTestCase, SkipIfBeforePyTorchVersion, TimedCall, skip_if_no_cuda, skip_if_quick
+
+
+def run_loading_test(num_workers=50, device=None, pw=False):
+    """multi workers stress tests"""
+    set_determinism(seed=0)
+    if device is None:
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    train_ds = list(range(10000))
+    train_loader = DataLoader(train_ds, batch_size=300, shuffle=True, num_workers=num_workers, persistent_workers=pw)
+    answer = []
+    for _ in range(2):
+        np.testing.assert_equal(torch.cuda.memory_allocated(), 0)
+        for batch_data in train_loader:
+            x = batch_data.to(device)
+            mem = torch.cuda.memory_allocated()
+            np.testing.assert_equal(mem > 0 and mem < 5000, True)
+        answer.append(x[-1].item())
+        del x
+    return answer
+
+
+@skip_if_quick
+@skip_if_no_cuda
+@SkipIfBeforePyTorchVersion((1, 9))
+class IntegrationLoading(DistTestCase):
+    def tearDown(self):
+        set_determinism(seed=None)
+
+    @TimedCall(seconds=5000, skip_timing=not torch.cuda.is_available(), daemon=False)
+    def test_timing(self):
+        expected = None
+        for pw in (False, True):
+            result = run_loading_test(pw=pw)
+            if expected is None:
+                expected = result[0]
+        np.testing.assert_allclose(result[0], expected)  # test for deterministic first epoch in two settings
+
+
+if __name__ == "__main__":
+    unittest.main()