def train(args):
# initialize Horovod library
hvd.init()
# Horovod limits CPU threads to be used per worker
torch.set_num_threads(1)
# disable logging for processes execpt 0 on every node
if hvd.local_rank() != 0:
f = open(os.devnull, "w")
sys.stdout = sys.stderr = f
elif not os.path.exists(args.dir):
# create 40 random image, mask paris on master node for training
print(
f"generating synthetic data to {args.dir} (this may take a while)")
os.makedirs(args.dir)
# set random seed to generate same random data for every node
np.random.seed(seed=0)
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(args.dir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(args.dir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(args.dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(args.dir, "seg*.nii.gz")))
train_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
train_transforms = Compose([
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys=["img", "seg"]),
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.5, spatial_axes=[0, 2]),
ToTensord(keys=["img", "seg"]),
])
# create a training data loader
train_ds = Dataset(data=train_files, transform=train_transforms)
# create a training data sampler
train_sampler = DistributedSampler(train_ds,
num_replicas=hvd.size(),
rank=hvd.rank())
# when supported, use "forkserver" to spawn dataloader workers instead of "fork" to prevent
# issues with Infiniband implementations that are not fork-safe
multiprocessing_context = None
if hasattr(
mp, "_supports_context"
) and mp._supports_context and "forkserver" in mp.get_all_start_methods():
multiprocessing_context = "forkserver"
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = DataLoader(
train_ds,
batch_size=2,
shuffle=False,
num_workers=1,
pin_memory=True,
sampler=train_sampler,
multiprocessing_context=multiprocessing_context,
)
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{hvd.local_rank()}")
model = monai.networks.nets.UNet(
dimensions=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).to(device)
optimizer = torch.optim.Adam(model.parameters(), 1e-3)
# Horovod broadcasts parameters & optimizer state
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)
# Horovod wraps optimizer with DistributedOptimizer
optimizer = hvd.DistributedOptimizer(
optimizer, named_parameters=model.named_parameters())
# start a typical PyTorch training
epoch_loss_values = list()
for epoch in range(5):
print("-" * 10)
print(f"epoch {epoch + 1}/{5}")
model.train()
epoch_loss = 0
step = 0
train_sampler.set_epoch(epoch)
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():.4f}")
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
print(f"train completed, epoch losses: {epoch_loss_values}")
if hvd.rank() == 0:
# all processes should see same parameters as they all start from same
# random parameters and gradients are synchronized in backward passes,
# therefore, saving it in one process is sufficient
torch.save(model.state_dict(), "final_model.pth")
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
tempdir = tempfile.mkdtemp()
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
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(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose([
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys=["img", "seg"]),
ToTensord(keys=["img", "seg"]),
])
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
# sliding window inference need to input 1 image in every iteration
val_loader = DataLoader(val_ds,
batch_size=1,
num_workers=4,
collate_fn=list_data_collate)
dice_metric = DiceMetric(include_background=True,
to_onehot_y=False,
sigmoid=True,
reduction="mean")
# try to use all the available GPUs
devices = get_devices_spec(None)
model = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(devices[0])
model.load_state_dict(torch.load("best_metric_model.pth"))
# if we have multiple GPUs, set data parallel to execute sliding window inference
if len(devices) > 1:
model = torch.nn.DataParallel(model, device_ids=devices)
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
saver = NiftiSaver(output_dir="./output")
for val_data in val_loader:
val_images, val_labels = val_data["img"].to(
devices[0]), val_data["seg"].to(devices[0])
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size,
sw_batch_size, model)
value = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
val_outputs = (val_outputs.sigmoid() >= 0.5).float()
saver.save_batch(val_outputs, val_data["img_meta_dict"])
metric = metric_sum / metric_count
print("evaluation metric:", metric)
shutil.rmtree(tempdir)
def evaluate(args):
if args.local_rank == 0 and not os.path.exists(args.dir):
# create 16 random image, mask paris for evaluation
print(
f"generating synthetic data to {args.dir} (this may take a while)")
os.makedirs(args.dir)
# set random seed to generate same random data for every node
np.random.seed(seed=0)
for i in range(16):
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(args.dir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(args.dir, f"seg{i:d}.nii.gz"))
# initialize the distributed evaluation process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
images = sorted(glob(os.path.join(args.dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(args.dir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose([
LoadImaged(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys="img"),
ToTensord(keys=["img", "seg"]),
])
# create a evaluation data loader
val_ds = Dataset(data=val_files, transform=val_transforms)
# create a evaluation data sampler
val_sampler = DistributedSampler(val_ds, shuffle=False)
# sliding window inference need to input 1 image in every iteration
val_loader = DataLoader(val_ds,
batch_size=1,
shuffle=False,
num_workers=2,
pin_memory=True,
sampler=val_sampler)
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_trans = Compose(
[Activations(sigmoid=True),
AsDiscrete(threshold_values=True)])
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
model = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
# wrap the model with DistributedDataParallel module
model = DistributedDataParallel(model, device_ids=[device])
# config mapping to expected GPU device
map_location = {"cuda:0": f"cuda:{args.local_rank}"}
# load model parameters to GPU device
model.load_state_dict(
torch.load("final_model.pth", map_location=map_location))
model.eval()
with torch.no_grad():
# define PyTorch Tensor to record metrics result at each GPU
# the first value is `sum` of all dice metric, the second value is `count` of not_nan items
metric = torch.zeros(2, dtype=torch.float, device=device)
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
roi_size = (96, 96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size,
sw_batch_size, model)
val_outputs = post_trans(val_outputs)
value, not_nans = dice_metric(y_pred=val_outputs, y=val_labels)
value = value.squeeze()
metric[0] += value * not_nans
metric[1] += not_nans
# synchronizes all processes and reduce results
dist.barrier()
dist.all_reduce(metric, op=torch.distributed.ReduceOp.SUM)
metric = metric.tolist()
if dist.get_rank() == 0:
print("evaluation metric:", metric[0] / metric[1])
dist.destroy_process_group()
def main(tempdir):
print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
im, _ = 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(tempdir, f"im{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
files = [{"img": img} for img in images]
# define pre transforms
pre_transforms = Compose([
LoadImaged(keys="img"),
EnsureChannelFirstd(keys="img"),
Orientationd(keys="img", axcodes="RAS"),
Resized(keys="img",
spatial_size=(96, 96, 96),
mode="trilinear",
align_corners=True),
ScaleIntensityd(keys="img"),
ToTensord(keys="img"),
])
# define dataset and dataloader
dataset = Dataset(data=files, transform=pre_transforms)
dataloader = DataLoader(dataset, batch_size=2, num_workers=4)
# define post transforms
post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
Invertd(
keys=
"pred", # invert the `pred` data field, also support multiple fields
transform=pre_transforms,
loader=dataloader,
orig_keys=
"img", # get the previously applied pre_transforms information on the `img` data field,
# then invert `pred` based on this information. we can use same info
# for multiple fields, also support different orig_keys for different fields
meta_keys=
"pred_meta_dict", # key field to save inverted meta data, every item maps to `keys`
orig_meta_keys=
"img_meta_dict", # get the meta data from `img_meta_dict` field when inverting,
# for example, may need the `affine` to invert `Spacingd` transform,
# multiple fields can use the same meta data to invert
meta_key_postfix=
"meta_dict", # if `meta_keys=None`, use "{keys}_{meta_key_postfix}" as the meta key,
# if `orig_meta_keys=None`, use "{orig_keys}_{meta_key_postfix}",
# otherwise, no need this arg during inverting
nearest_interp=
True, # change to use "nearest" mode in interpolation when inverting
to_tensor=True, # convert to PyTorch Tensor after inverting
),
SaveImaged(keys="pred",
meta_keys="pred_meta_dict",
output_dir="./out",
output_postfix="seg",
resample=False),
])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
net.load_state_dict(
torch.load("best_metric_model_segmentation3d_dict.pth"))
net.eval()
with torch.no_grad():
for d in dataloader:
images = d["img"].to(device)
# define sliding window size and batch size for windows inference
d["pred"] = sliding_window_inference(inputs=images,
roi_size=(96, 96, 96),
sw_batch_size=4,
predictor=net)
# execute post transforms to invert spatial transforms and save to NIfTI files
post_transforms(d)
def main(tempdir):
config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
# define transforms for image and segmentation
imtrans = Compose([ScaleIntensity(), AddChannel(), ToTensor()])
segtrans = Compose([AddChannel(), ToTensor()])
ds = NiftiDataset(images,
segs,
transform=imtrans,
seg_transform=segtrans,
image_only=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
post_trans = Compose(
[Activations(sigmoid=True),
AsDiscrete(threshold_values=True)])
def _sliding_window_processor(engine, batch):
net.eval()
with torch.no_grad():
val_images, val_labels = batch[0].to(device), batch[1].to(device)
seg_probs = sliding_window_inference(val_images, roi_size,
sw_batch_size, net)
seg_probs = post_trans(seg_probs)
return seg_probs, val_labels
evaluator = Engine(_sliding_window_processor)
# add evaluation metric to the evaluator engine
MeanDice().attach(evaluator, "Mean_Dice")
# StatsHandler prints loss at every iteration and print metrics at every epoch,
# we don't need to print loss for evaluator, so just print metrics, user can also customize print functions
val_stats_handler = StatsHandler(
name="evaluator",
output_transform=lambda x:
None, # no need to print loss value, so disable per iteration output
)
val_stats_handler.attach(evaluator)
# for the array data format, assume the 3rd item of batch data is the meta_data
file_saver = SegmentationSaver(
output_dir="tempdir",
output_ext=".nii.gz",
output_postfix="seg",
name="evaluator",
batch_transform=lambda x: x[2],
output_transform=lambda output: output[0],
)
file_saver.attach(evaluator)
# the model was trained by "unet_training_array" example
ckpt_saver = CheckpointLoader(
load_path="./runs_array/net_checkpoint_100.pt", load_dict={"net": net})
ckpt_saver.attach(evaluator)
# sliding window inference for one image at every iteration
loader = DataLoader(ds,
batch_size=1,
num_workers=1,
pin_memory=torch.cuda.is_available())
state = evaluator.run(loader)
print(state)
def test_invert(self):
set_determinism(seed=0)
im_fname, seg_fname = [
make_nifti_image(i)
for i in create_test_image_3d(101, 100, 107, noise_max=100)
]
transform = Compose([
LoadImaged(KEYS),
AddChanneld(KEYS),
Orientationd(KEYS, "RPS"),
Spacingd(KEYS,
pixdim=(1.2, 1.01, 0.9),
mode=["bilinear", "nearest"],
dtype=np.float32),
ScaleIntensityd("image", minv=1, maxv=10),
RandFlipd(KEYS, prob=0.5, spatial_axis=[1, 2]),
RandAxisFlipd(KEYS, prob=0.5),
RandRotate90d(KEYS, spatial_axes=(1, 2)),
RandZoomd(KEYS,
prob=0.5,
min_zoom=0.5,
max_zoom=1.1,
keep_size=True),
RandRotated(KEYS,
prob=0.5,
range_x=np.pi,
mode="bilinear",
align_corners=True),
RandAffined(KEYS, prob=0.5, rotate_range=np.pi, mode="nearest"),
ResizeWithPadOrCropd(KEYS, 100),
ToTensord(
"image"
), # test to support both Tensor and Numpy array when inverting
CastToTyped(KEYS, dtype=[torch.uint8, np.uint8]),
])
data = [{"image": im_fname, "label": seg_fname} for _ in range(12)]
# num workers = 0 for mac or gpu transforms
num_workers = 0 if sys.platform == "darwin" or torch.cuda.is_available(
) else 2
dataset = CacheDataset(data, transform=transform, progress=False)
loader = DataLoader(dataset, num_workers=num_workers, batch_size=5)
# set up engine
def _train_func(engine, batch):
self.assertTupleEqual(batch["image"].shape[1:], (1, 100, 100, 100))
engine.state.output = batch
engine.fire_event(IterationEvents.MODEL_COMPLETED)
return engine.state.output
engine = Engine(_train_func)
engine.register_events(*IterationEvents)
# set up testing handler
TransformInverter(
transform=transform,
loader=loader,
output_keys=["image", "label"],
batch_keys="label",
nearest_interp=True,
postfix="inverted1",
to_tensor=[True, False],
device="cpu",
num_workers=0
if sys.platform == "darwin" or torch.cuda.is_available() else 2,
).attach(engine)
# test different nearest interpolation values
TransformInverter(
transform=transform,
loader=loader,
output_keys=["image", "label"],
batch_keys="image",
nearest_interp=[True, False],
post_func=[lambda x: x + 10, lambda x: x],
postfix="inverted2",
collate_fn=pad_list_data_collate,
num_workers=0
if sys.platform == "darwin" or torch.cuda.is_available() else 2,
).attach(engine)
engine.run(loader, max_epochs=1)
set_determinism(seed=None)
self.assertTupleEqual(engine.state.output["image"].shape,
(2, 1, 100, 100, 100))
self.assertTupleEqual(engine.state.output["label"].shape,
(2, 1, 100, 100, 100))
# check the nearest inerpolation mode
for i in engine.state.output["image_inverted1"]:
torch.testing.assert_allclose(
i.to(torch.uint8).to(torch.float), i.to(torch.float))
self.assertTupleEqual(i.shape, (1, 100, 101, 107))
for i in engine.state.output["label_inverted1"]:
np.testing.assert_allclose(
i.astype(np.uint8).astype(np.float32), i.astype(np.float32))
self.assertTupleEqual(i.shape, (1, 100, 101, 107))
# check labels match
reverted = engine.state.output["label_inverted1"][-1].astype(np.int32)
original = LoadImaged(KEYS)(data[-1])["label"]
n_good = np.sum(np.isclose(reverted, original, atol=1e-3))
reverted_name = engine.state.output["label_meta_dict"][
"filename_or_obj"][-1]
original_name = data[-1]["label"]
self.assertEqual(reverted_name, original_name)
print("invert diff", reverted.size - n_good)
# 25300: 2 workers (cpu, non-macos)
# 1812: 0 workers (gpu or macos)
# 1824: torch 1.5.1
self.assertTrue((reverted.size - n_good) in (25300, 1812, 1824),
"diff. in 3 possible values")
# check the case that different items use different interpolation mode to invert transforms
d = engine.state.output["image_inverted2"]
# if the interpolation mode is nearest, accumulated diff should be smaller than 1
self.assertLess(
torch.sum(d.to(torch.float) -
d.to(torch.uint8).to(torch.float)).item(), 1.0)
self.assertTupleEqual(d.shape, (2, 1, 100, 101, 107))
d = engine.state.output["label_inverted2"]
# if the interpolation mode is not nearest, accumulated diff should be greater than 10000
self.assertGreater(
torch.sum(d.to(torch.float) -
d.to(torch.uint8).to(torch.float)).item(), 10000.0)
self.assertTupleEqual(d.shape, (2, 1, 100, 101, 107))
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
tempdir = tempfile.mkdtemp()
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
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(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose(
[
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys=["img", "seg"]),
ToTensord(keys=["img", "seg"]),
]
)
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
device = torch.device("cuda:0")
net = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
)
net.to(device)
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
def _sliding_window_processor(engine, batch):
net.eval()
with torch.no_grad():
val_images, val_labels = batch["img"].to(device), batch["seg"].to(device)
seg_probs = sliding_window_inference(val_images, roi_size, sw_batch_size, net)
return seg_probs, val_labels
evaluator = Engine(_sliding_window_processor)
# add evaluation metric to the evaluator engine
MeanDice(sigmoid=True, to_onehot_y=False).attach(evaluator, "Mean_Dice")
# StatsHandler prints loss at every iteration and print metrics at every epoch,
# we don't need to print loss for evaluator, so just print metrics, user can also customize print functions
val_stats_handler = StatsHandler(
name="evaluator",
output_transform=lambda x: None, # no need to print loss value, so disable per iteration output
)
val_stats_handler.attach(evaluator)
# convert the necessary metadata from batch data
SegmentationSaver(
output_dir="tempdir",
output_ext=".nii.gz",
output_postfix="seg",
name="evaluator",
batch_transform=lambda batch: batch["img_meta_dict"],
output_transform=lambda output: predict_segmentation(output[0]),
).attach(evaluator)
# the model was trained by "unet_training_dict" example
CheckpointLoader(load_path="./runs/net_checkpoint_50.pth", load_dict={"net": net}).attach(evaluator)
# sliding window inference for one image at every iteration
val_loader = DataLoader(
val_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate, pin_memory=torch.cuda.is_available()
)
state = evaluator.run(val_loader)
print(state)
shutil.rmtree(tempdir)
def main(tempdir):
config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
# define transforms for image and segmentation
imtrans = Compose([ScaleIntensity(), AddChannel(), ToTensor()])
segtrans = Compose([AddChannel(), ToTensor()])
val_ds = NiftiDataset(images,
segs,
transform=imtrans,
seg_transform=segtrans,
image_only=False)
# sliding window inference for one image at every iteration
val_loader = DataLoader(val_ds,
batch_size=1,
num_workers=1,
pin_memory=torch.cuda.is_available())
dice_metric = DiceMetric(include_background=True,
to_onehot_y=False,
sigmoid=True,
reduction="mean")
device = torch.device("cuda:0")
model = UNet(
dimensions=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.load_state_dict(torch.load("best_metric_model.pth"))
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
saver = NiftiSaver(output_dir="./output")
for val_data in val_loader:
val_images, val_labels = val_data[0].to(device), val_data[1].to(
device)
# define sliding window size and batch size for windows inference
roi_size = (96, 96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size,
sw_batch_size, model)
value = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
val_outputs = (val_outputs.sigmoid() >= 0.5).float()
saver.save_batch(val_outputs, val_data[2])
metric = metric_sum / metric_count
print("evaluation metric:", metric)
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask paris
tempdir = tempfile.mkdtemp()
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
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(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
val_files = [{
"image": img,
"label": seg
} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys=["image", "label"]),
ToTensord(keys=["image", "label"]),
])
# 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)
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda:0")
net = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
val_post_transforms = Compose([
Activationsd(keys="pred", output_postfix="act", sigmoid=True),
AsDiscreted(keys="pred_act",
output_postfix="dis",
threshold_values=True),
KeepLargestConnectedComponentd(keys="pred_act_dis",
applied_values=[1],
output_postfix=None),
])
val_handlers = [
StatsHandler(output_transform=lambda x: None),
CheckpointLoader(load_path="./runs/net_key_metric=0.9101.pth",
load_dict={"net": net}),
SegmentationSaver(
output_dir="./runs/",
batch_transform=lambda x: {
"filename_or_obj": x["image.filename_or_obj"],
"affine": x["image.affine"],
"original_affine": x["image.original_affine"],
"spatial_shape": x["image.spatial_shape"],
},
output_transform=lambda x: x["pred_act_dis"],
),
]
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=SlidingWindowInferer(roi_size=(96, 96, 96),
sw_batch_size=4,
overlap=0.5),
post_transform=val_post_transforms,
key_val_metric={
"val_mean_dice":
MeanDice(include_background=True,
output_transform=lambda x:
(x["pred_act_dis"], x["label"]))
},
additional_metrics={
"val_acc":
Accuracy(
output_transform=lambda x: (x["pred_act_dis"], x["label"]))
},
val_handlers=val_handlers,
)
evaluator.run()
shutil.rmtree(tempdir)
def main(tempdir):
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask pairs
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(5):
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(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
val_files = [{"image": img, "label": seg} for img, seg in zip(images, segs)]
# model file path
model_file = glob("./runs/net_key_metric*")[0]
# define transforms for image and segmentation
val_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys="image"),
ToTensord(keys=["image", "label"]),
]
)
# 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)
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
val_post_transforms = Compose(
[
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
]
)
val_handlers = [
StatsHandler(output_transform=lambda x: None),
CheckpointLoader(load_path=model_file, load_dict={"net": net}),
SegmentationSaver(
output_dir="./runs/",
batch_transform=lambda batch: batch["image_meta_dict"],
output_transform=lambda output: output["pred"],
),
]
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=SlidingWindowInferer(roi_size=(96, 96, 96), sw_batch_size=4, overlap=0.5),
post_transform=val_post_transforms,
key_val_metric={
"val_mean_dice": MeanDice(include_background=True, output_transform=lambda x: (x["pred"], x["label"]))
},
additional_metrics={"val_acc": Accuracy(output_transform=lambda x: (x["pred"], x["label"]))},
val_handlers=val_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP evaluation
amp=True if monai.config.get_torch_version_tuple() >= (1, 6) else False,
)
evaluator.run()
def main(tempdir):
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask pairs
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(40):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
# define transforms for image and segmentation
train_imtrans = Compose([
ScaleIntensity(),
AddChannel(),
RandSpatialCrop((96, 96, 96), random_size=False),
EnsureType(),
])
train_segtrans = Compose([
AddChannel(),
RandSpatialCrop((96, 96, 96), random_size=False),
EnsureType()
])
val_imtrans = Compose(
[ScaleIntensity(),
AddChannel(),
Resize((96, 96, 96)),
EnsureType()])
val_segtrans = Compose([AddChannel(), Resize((96, 96, 96)), EnsureType()])
# define image dataset, data loader
check_ds = ImageDataset(images,
segs,
transform=train_imtrans,
seg_transform=train_segtrans)
check_loader = DataLoader(check_ds,
batch_size=10,
num_workers=2,
pin_memory=torch.cuda.is_available())
im, seg = monai.utils.misc.first(check_loader)
print(im.shape, seg.shape)
# create a training data loader
train_ds = ImageDataset(images[:20],
segs[:20],
transform=train_imtrans,
seg_transform=train_segtrans)
train_loader = DataLoader(
train_ds,
batch_size=5,
shuffle=True,
num_workers=8,
pin_memory=torch.cuda.is_available(),
)
# create a validation data loader
val_ds = ImageDataset(images[-20:],
segs[-20:],
transform=val_imtrans,
seg_transform=val_segtrans)
val_loader = DataLoader(val_ds,
batch_size=5,
num_workers=8,
pin_memory=torch.cuda.is_available())
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = 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 = monai.losses.DiceLoss(sigmoid=True)
lr = 1e-3
opt = torch.optim.Adam(net.parameters(), lr)
# Ignite trainer expects batch=(img, seg) and returns output=loss at every iteration,
# user can add output_transform to return other values, like: y_pred, y, etc.
trainer = create_supervised_trainer(net, opt, loss, device, False)
# adding checkpoint handler to save models (network params and optimizer stats) during training
checkpoint_handler = ModelCheckpoint("./runs_array/",
"net",
n_saved=10,
require_empty=False)
trainer.add_event_handler(
event_name=Events.EPOCH_COMPLETED,
handler=checkpoint_handler,
to_save={
"net": net,
"opt": opt
},
)
# StatsHandler prints loss at every iteration and print metrics at every epoch,
# we don't set metrics for trainer here, so just print loss, user can also customize print functions
# and can use output_transform to convert engine.state.output if it's not a loss value
train_stats_handler = StatsHandler(name="trainer",
output_transform=lambda x: x)
train_stats_handler.attach(trainer)
# TensorBoardStatsHandler plots loss at every iteration and plots metrics at every epoch, same as StatsHandler
train_tensorboard_stats_handler = TensorBoardStatsHandler(
output_transform=lambda x: x)
train_tensorboard_stats_handler.attach(trainer)
validation_every_n_epochs = 1
# Set parameters for validation
metric_name = "Mean_Dice"
# add evaluation metric to the evaluator engine
val_metrics = {metric_name: MeanDice()}
post_pred = Compose(
[EnsureType(),
Activations(sigmoid=True),
AsDiscrete(threshold=0.5)])
post_label = Compose([EnsureType(), AsDiscrete(threshold=0.5)])
# Ignite evaluator expects batch=(img, seg) and returns output=(y_pred, y) at every iteration,
# user can add output_transform to return other values
evaluator = create_supervised_evaluator(
net,
val_metrics,
device,
True,
output_transform=lambda x, y, y_pred:
([post_pred(i) for i in decollate_batch(y_pred)],
[post_label(i) for i in decollate_batch(y)]),
)
@trainer.on(Events.EPOCH_COMPLETED(every=validation_every_n_epochs))
def run_validation(engine):
evaluator.run(val_loader)
# add early stopping handler to evaluator
early_stopper = EarlyStopping(
patience=4,
score_function=stopping_fn_from_metric(metric_name),
trainer=trainer)
evaluator.add_event_handler(event_name=Events.EPOCH_COMPLETED,
handler=early_stopper)
# add stats event handler to print validation stats via evaluator
val_stats_handler = StatsHandler(
name="evaluator",
output_transform=lambda x:
None, # no need to print loss value, so disable per iteration output
global_epoch_transform=lambda x: trainer.state.epoch,
) # fetch global epoch number from trainer
val_stats_handler.attach(evaluator)
# add handler to record metrics to TensorBoard at every validation epoch
val_tensorboard_stats_handler = TensorBoardStatsHandler(
output_transform=lambda x:
None, # no need to plot loss value, so disable per iteration output
global_epoch_transform=lambda x: trainer.state.epoch,
) # fetch global epoch number from trainer
val_tensorboard_stats_handler.attach(evaluator)
# add handler to draw the first image and the corresponding label and model output in the last batch
# here we draw the 3D output as GIF format along Depth axis, at every validation epoch
val_tensorboard_image_handler = TensorBoardImageHandler(
batch_transform=lambda batch: (batch[0], batch[1]),
output_transform=lambda output: output[0],
global_iter_transform=lambda x: trainer.state.epoch,
)
evaluator.add_event_handler(event_name=Events.EPOCH_COMPLETED,
handler=val_tensorboard_image_handler)
train_epochs = 30
state = trainer.run(train_loader, train_epochs)
print(state)
def train(args):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
if idist.get_local_rank() == 0 and not os.path.exists(args.dir):
# create 40 random image, mask paris for training
print(f"generating synthetic data to {args.dir} (this may take a while)")
os.makedirs(args.dir)
# set random seed to generate same random data for every node
np.random.seed(seed=0)
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(args.dir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(args.dir, f"seg{i:d}.nii.gz"))
idist.barrier()
images = sorted(glob(os.path.join(args.dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(args.dir, "seg*.nii.gz")))
train_files = [{"image": img, "label": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
train_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys="image"),
RandCropByPosNegLabeld(
keys=["image", "label"], label_key="label", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
),
RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]),
EnsureTyped(keys=["image", "label"]),
]
)
# create a training data loader
train_ds = Dataset(data=train_files, transform=train_transforms)
# create a training data sampler
train_sampler = DistributedSampler(train_ds)
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = DataLoader(
train_ds,
batch_size=2,
shuffle=False,
num_workers=2,
pin_memory=True,
sampler=train_sampler,
)
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{idist.get_local_rank()}")
torch.cuda.set_device(device)
net = 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 = monai.losses.DiceLoss(sigmoid=True)
opt = torch.optim.Adam(net.parameters(), 1e-3)
lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1)
# wrap the model with DistributedDataParallel module
net = DistributedDataParallel(net, device_ids=[device])
train_post_transforms = Compose(
[
EnsureTyped(keys="pred"),
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold=0.5),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
]
)
train_handlers = [
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
]
if idist.get_rank() == 0:
train_handlers.extend(
[
StatsHandler(tag_name="train_loss", output_transform=from_engine(["loss"], first=True)),
CheckpointSaver(save_dir="./runs/", save_dict={"net": net, "opt": opt}, save_interval=2),
]
)
trainer = SupervisedTrainer(
device=device,
max_epochs=5,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=loss,
inferer=SimpleInferer(),
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP evaluation
amp=True if monai.utils.get_torch_version_tuple() >= (1, 6) else False,
postprocessing=train_post_transforms,
key_train_metric={"train_acc": Accuracy(output_transform=from_engine(["pred", "label"]), device=device)},
train_handlers=train_handlers,
)
trainer.run()
dist.destroy_process_group()
import monai
from monai.data import NiftiDataset, create_test_image_3d
from monai.transforms import Compose, AddChannel, ScaleIntensity, RandSpatialCrop, Resize, ToTensor
from monai.handlers import \
StatsHandler, TensorBoardStatsHandler, TensorBoardImageHandler, MeanDice, stopping_fn_from_metric
from monai.networks.utils import predict_segmentation
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask paris
tempdir = tempfile.mkdtemp()
print(
'generating synthetic data to {} (this may take a while)'.format(tempdir))
for i in range(40):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, 'im%i.nii.gz' % i))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, 'seg%i.nii.gz' % i))
images = sorted(glob(os.path.join(tempdir, 'im*.nii.gz')))
segs = sorted(glob(os.path.join(tempdir, 'seg*.nii.gz')))
# define transforms for image and segmentation
train_imtrans = Compose([
ScaleIntensity(),
AddChannel(),
RandSpatialCrop((96, 96, 96), random_size=False),
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask paris
tempdir = tempfile.mkdtemp()
print(f"generating synthetic data to {tempdir} (this may take a while)")
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(tempdir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
train_files = [{Keys.IMAGE: img, Keys.LABEL: seg} for img, seg in zip(images[:20], segs[:20])]
val_files = [{Keys.IMAGE: img, Keys.LABEL: seg} for img, seg in zip(images[-20:], segs[-20:])]
# define transforms for image and segmentation
train_transforms = Compose(
[
LoadNiftid(keys=[Keys.IMAGE, Keys.LABEL]),
AsChannelFirstd(keys=[Keys.IMAGE, Keys.LABEL], channel_dim=-1),
ScaleIntensityd(keys=[Keys.IMAGE, Keys.LABEL]),
RandCropByPosNegLabeld(
keys=[Keys.IMAGE, Keys.LABEL], label_key=Keys.LABEL, size=[96, 96, 96], pos=1, neg=1, num_samples=4
),
RandRotate90d(keys=[Keys.IMAGE, Keys.LABEL], prob=0.5, spatial_axes=[0, 2]),
ToTensord(keys=[Keys.IMAGE, Keys.LABEL]),
]
)
val_transforms = Compose(
[
LoadNiftid(keys=[Keys.IMAGE, Keys.LABEL]),
AsChannelFirstd(keys=[Keys.IMAGE, Keys.LABEL], channel_dim=-1),
ScaleIntensityd(keys=[Keys.IMAGE, Keys.LABEL]),
ToTensord(keys=[Keys.IMAGE, Keys.LABEL]),
]
)
# create a training data loader
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 = DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4, collate_fn=list_data_collate)
# create a validation data loader
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate)
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda:0")
net = monai.networks.nets.UNet(
dimensions=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 = monai.losses.DiceLoss(do_sigmoid=True)
opt = torch.optim.Adam(net.parameters(), 1e-3)
val_post_transforms = Compose(
[
Activationsd(keys=Keys.PRED, output_postfix="act", sigmoid=True),
AsDiscreted(keys="pred_act", output_postfix="dis", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred_act_dis", applied_values=[1], output_postfix=None),
]
)
val_handlers = [StatsHandler(output_transform=lambda x: None)]
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=SlidingWindowInferer(roi_size=(96, 96, 96), sw_batch_size=4, overlap=0.5),
post_transform=val_post_transforms,
key_val_metric={
"val_mean_dice": MeanDice(
include_background=True, output_transform=lambda x: (x["pred_act_dis"], x[Keys.LABEL])
)
},
additional_metrics={"val_acc": Accuracy(output_transform=lambda x: (x["pred_act_dis"], x[Keys.LABEL]))},
val_handlers=val_handlers,
)
train_post_transforms = Compose(
[
Activationsd(keys=Keys.PRED, output_postfix="act", sigmoid=True),
AsDiscreted(keys="pred_act", output_postfix="dis", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred_act_dis", applied_values=[1], output_postfix=None),
]
)
train_handlers = [
ValidationHandler(validator=evaluator, interval=2, epoch_level=True),
StatsHandler(tag_name="train_loss", output_transform=lambda x: x[Keys.INFO][Keys.LOSS]),
]
trainer = SupervisedTrainer(
device=device,
max_epochs=5,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=loss,
inferer=SimpleInferer(),
amp=False,
post_transform=train_post_transforms,
key_train_metric={"train_acc": Accuracy(output_transform=lambda x: (x["pred_act_dis"], x[Keys.LABEL]))},
train_handlers=train_handlers,
)
trainer.run()
shutil.rmtree(tempdir)
def setUp(self):
im, msk = create_test_image_3d(self.im_shape[0], self.im_shape[1], self.im_shape[2], 4, 20, 0, self.num_classes)
self.imt = im[None, None]
self.seg1 = (msk[None, None] > 0).astype(np.float32)
self.segn = msk[None, None]
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask paris
tempdir = tempfile.mkdtemp()
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(40):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"im{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
# define transforms for image and segmentation
train_imtrans = Compose([
ScaleIntensity(),
AddChannel(),
RandSpatialCrop((96, 96, 96), random_size=False),
RandRotate90(prob=0.5, spatial_axes=(0, 2)),
ToTensor(),
])
train_segtrans = Compose([
AddChannel(),
RandSpatialCrop((96, 96, 96), random_size=False),
RandRotate90(prob=0.5, spatial_axes=(0, 2)),
ToTensor(),
])
val_imtrans = Compose([ScaleIntensity(), AddChannel(), ToTensor()])
val_segtrans = Compose([AddChannel(), ToTensor()])
# define nifti dataset, data loader
check_ds = NiftiDataset(images,
segs,
transform=train_imtrans,
seg_transform=train_segtrans)
check_loader = DataLoader(check_ds,
batch_size=10,
num_workers=2,
pin_memory=torch.cuda.is_available())
im, seg = monai.utils.misc.first(check_loader)
print(im.shape, seg.shape)
# create a training data loader
train_ds = NiftiDataset(images[:20],
segs[:20],
transform=train_imtrans,
seg_transform=train_segtrans)
train_loader = DataLoader(train_ds,
batch_size=4,
shuffle=True,
num_workers=8,
pin_memory=torch.cuda.is_available())
# create a validation data loader
val_ds = NiftiDataset(images[-20:],
segs[-20:],
transform=val_imtrans,
seg_transform=val_segtrans)
val_loader = DataLoader(val_ds,
batch_size=1,
num_workers=4,
pin_memory=torch.cuda.is_available())
dice_metric = DiceMetric(include_background=True,
to_onehot_y=False,
sigmoid=True,
reduction="mean")
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda:0")
model = monai.networks.nets.UNet(
dimensions=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(), 1e-3)
# start a typical PyTorch training
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter()
for epoch in range(5):
print("-" * 10)
print(f"epoch {epoch + 1}/{5}")
model.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs, labels = batch_data[0].to(device), batch_data[1].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():.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:.4f}")
if (epoch + 1) % val_interval == 0:
model.eval()
with torch.no_grad():
metric_sum = 0.0
metric_count = 0
val_images = None
val_labels = None
val_outputs = None
for val_data in val_loader:
val_images, val_labels = val_data[0].to(
device), val_data[1].to(device)
roi_size = (96, 96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(
val_images, roi_size, sw_batch_size, model)
value = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
metric = metric_sum / metric_count
metric_values.append(metric)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
torch.save(model.state_dict(), "best_metric_model.pth")
print("saved new best metric model")
print(
"current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}"
.format(epoch + 1, metric, best_metric, 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")
shutil.rmtree(tempdir)
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
writer.close()
def train(args):
# disable logging for processes except 0 on every node
if args.local_rank != 0:
f = open(os.devnull, "w")
sys.stdout = sys.stderr = f
elif not os.path.exists(args.dir):
# create 40 random image, mask paris for training
print(
f"generating synthetic data to {args.dir} (this may take a while)")
os.makedirs(args.dir)
# set random seed to generate same random data for every node
np.random.seed(seed=0)
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(args.dir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(args.dir, f"seg{i:d}.nii.gz"))
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
images = sorted(glob(os.path.join(args.dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(args.dir, "seg*.nii.gz")))
train_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
train_transforms = Compose([
LoadImaged(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
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.5, spatial_axes=[0, 2]),
EnsureTyped(keys=["img", "seg"]),
])
# create a training data loader
train_ds = Dataset(data=train_files, transform=train_transforms)
# create a training data sampler
train_sampler = DistributedSampler(dataset=train_ds,
even_divisible=True,
shuffle=True)
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = DataLoader(
train_ds,
batch_size=2,
shuffle=False,
num_workers=2,
pin_memory=True,
sampler=train_sampler,
)
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
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).to(device)
optimizer = torch.optim.Adam(model.parameters(), 1e-3)
# wrap the model with DistributedDataParallel module
model = DistributedDataParallel(model, device_ids=[device])
# start a typical PyTorch training
epoch_loss_values = list()
for epoch in range(5):
print("-" * 10)
print(f"epoch {epoch + 1}/{5}")
model.train()
epoch_loss = 0
step = 0
train_sampler.set_epoch(epoch)
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():.4f}")
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
print(f"train completed, epoch losses: {epoch_loss_values}")
if dist.get_rank() == 0:
# all processes should see same parameters as they all start from same
# random parameters and gradients are synchronized in backward passes,
# therefore, saving it in one process is sufficient
torch.save(model.state_dict(), "final_model.pth")
dist.destroy_process_group()
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask paris
tempdir = tempfile.mkdtemp()
print(f"generating synthetic data to {tempdir} (this may take a while)")
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(tempdir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
train_files = [{"image": img, "label": seg} for img, seg in zip(images[:20], segs[:20])]
val_files = [{"image": img, "label": seg} for img, seg in zip(images[-20:], segs[-20:])]
# define transforms for image and segmentation
train_transforms = Compose(
[
LoadNiftid(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys=["image", "label"]),
RandCropByPosNegLabeld(
keys=["image", "label"], label_key="label", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
),
RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]),
ToTensord(keys=["image", "label"]),
]
)
val_transforms = Compose(
[
LoadNiftid(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys=["image", "label"]),
ToTensord(keys=["image", "label"]),
]
)
# create a training data loader
train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.5)
# 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.CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0)
val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda:0")
net = monai.networks.nets.UNet(
dimensions=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 = monai.losses.DiceLoss(sigmoid=True)
opt = torch.optim.Adam(net.parameters(), 1e-3)
lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1)
val_post_transforms = Compose(
[
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
]
)
val_handlers = [
StatsHandler(output_transform=lambda x: None),
TensorBoardStatsHandler(log_dir="./runs/", output_transform=lambda x: None),
TensorBoardImageHandler(
log_dir="./runs/", batch_transform=lambda x: (x["image"], x["label"]), output_transform=lambda x: x["pred"]
),
CheckpointSaver(save_dir="./runs/", save_dict={"net": net}, save_key_metric=True),
]
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=SlidingWindowInferer(roi_size=(96, 96, 96), sw_batch_size=4, overlap=0.5),
post_transform=val_post_transforms,
key_val_metric={
"val_mean_dice": MeanDice(include_background=True, output_transform=lambda x: (x["pred"], x["label"]))
},
additional_metrics={"val_acc": Accuracy(output_transform=lambda x: (x["pred"], x["label"]))},
val_handlers=val_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP evaluation
amp=True if monai.config.get_torch_version_tuple() >= (1, 6) else False,
)
train_post_transforms = Compose(
[
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
]
)
train_handlers = [
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
ValidationHandler(validator=evaluator, interval=2, epoch_level=True),
StatsHandler(tag_name="train_loss", output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(log_dir="./runs/", tag_name="train_loss", output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir="./runs/", save_dict={"net": net, "opt": opt}, save_interval=2, epoch_level=True),
]
trainer = SupervisedTrainer(
device=device,
max_epochs=5,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=loss,
inferer=SimpleInferer(),
post_transform=train_post_transforms,
key_train_metric={"train_acc": Accuracy(output_transform=lambda x: (x["pred"], x["label"]))},
train_handlers=train_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP training
amp=True if monai.config.get_torch_version_tuple() >= (1, 6) else False,
)
trainer.run()
shutil.rmtree(tempdir)
def evaluate(args):
# initialize Horovod library
hvd.init()
# Horovod limits CPU threads to be used per worker
torch.set_num_threads(1)
if hvd.local_rank() == 0 and not os.path.exists(args.dir):
# create 16 random image, mask paris for evaluation
print(
f"generating synthetic data to {args.dir} (this may take a while)")
os.makedirs(args.dir)
# set random seed to generate same random data for every node
np.random.seed(seed=0)
for i in range(16):
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(args.dir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(args.dir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(args.dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(args.dir, "seg*.nii.gz")))
val_files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose([
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys="img"),
ToTensord(keys=["img", "seg"]),
])
# create a evaluation data loader
val_ds = Dataset(data=val_files, transform=val_transforms)
# create a evaluation data sampler
val_sampler = DistributedSampler(val_ds,
shuffle=False,
num_replicas=hvd.size(),
rank=hvd.rank())
# when supported, use "forkserver" to spawn dataloader workers instead of "fork" to prevent
# issues with Infiniband implementations that are not fork-safe
multiprocessing_context = None
if hasattr(
mp, "_supports_context"
) and mp._supports_context and "forkserver" in mp.get_all_start_methods():
multiprocessing_context = "forkserver"
# sliding window inference need to input 1 image in every iteration
val_loader = DataLoader(
val_ds,
batch_size=1,
shuffle=False,
num_workers=2,
pin_memory=True,
sampler=val_sampler,
multiprocessing_context=multiprocessing_context,
)
dice_metric = DiceMetric(include_background=True,
to_onehot_y=False,
sigmoid=True,
reduction="mean")
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{hvd.local_rank()}")
torch.cuda.set_device(device)
model = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
if hvd.rank() == 0:
# load model parameters for evaluation
model.load_state_dict(torch.load("final_model.pth"))
# Horovod broadcasts parameters
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
model.eval()
with torch.no_grad():
# define PyTorch Tensor to record metrics result at each GPU
# the first value is `sum` of all dice metric, the second value is `count` of not_nan items
metric = torch.zeros(2, dtype=torch.float, device=device)
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
roi_size = (96, 96, 96)
sw_batch_size = 4
val_outputs = sliding_window_inference(val_images, roi_size,
sw_batch_size, model)
value = dice_metric(y_pred=val_outputs, y=val_labels).squeeze()
metric[0] += value * dice_metric.not_nans
metric[1] += dice_metric.not_nans
# synchronizes all processes and reduce results
print(
f"metric in rank {hvd.rank()}: sum={metric[0].item()}, count={metric[1].item()}"
)
avg_metric = hvd.allreduce(metric, name="mean_dice")
if hvd.rank() == 0:
print(
f"average metric: sum={avg_metric[0].item()}, count={avg_metric[1].item()}"
)
print("evaluation metric:", (avg_metric[0] / avg_metric[1]).item())
def evaluate(args):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
if args.local_rank == 0 and not os.path.exists(args.dir):
# create 16 random image, mask paris for evaluation
print(f"generating synthetic data to {args.dir} (this may take a while)")
os.makedirs(args.dir)
# set random seed to generate same random data for every node
np.random.seed(seed=0)
for i in range(16):
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(args.dir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(args.dir, f"seg{i:d}.nii.gz"))
# initialize the distributed evaluation process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
images = sorted(glob(os.path.join(args.dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(args.dir, "seg*.nii.gz")))
val_files = [{"image": img, "label": seg} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
val_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys="image"),
ToTensord(keys=["image", "label"]),
]
)
# create a evaluation data loader
val_ds = Dataset(data=val_files, transform=val_transforms)
# create a evaluation data sampler
val_sampler = DistributedSampler(val_ds, shuffle=False)
# sliding window inference need to input 1 image in every iteration
val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=2, pin_memory=True, sampler=val_sampler)
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
net = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
# wrap the model with DistributedDataParallel module
net = DistributedDataParallel(net, device_ids=[device])
val_post_transforms = Compose(
[
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
]
)
val_handlers = [
CheckpointLoader(
load_path="./runs/checkpoint_epoch=4.pt",
load_dict={"net": net},
# config mapping to expected GPU device
map_location={"cuda:0": f"cuda:{args.local_rank}"},
),
]
if dist.get_rank() == 0:
val_handlers.extend(
[
StatsHandler(output_transform=lambda x: None),
SegmentationSaver(
output_dir="./runs/",
batch_transform=lambda batch: batch["image_meta_dict"],
output_transform=lambda output: output["pred"],
),
]
)
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=SlidingWindowInferer(roi_size=(96, 96, 96), sw_batch_size=4, overlap=0.5),
post_transform=val_post_transforms,
key_val_metric={
"val_mean_dice": MeanDice(
include_background=True,
output_transform=lambda x: (x["pred"], x["label"]),
device=device,
)
},
additional_metrics={"val_acc": Accuracy(output_transform=lambda x: (x["pred"], x["label"]), device=device)},
val_handlers=val_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP evaluation
amp=True if monai.utils.get_torch_version_tuple() >= (1, 6) else False,
)
evaluator.run()
dist.destroy_process_group()
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# create a temporary directory and 40 random image, mask paris
tempdir = tempfile.mkdtemp()
print(f"generating synthetic data to {tempdir} (this may take a while)")
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(tempdir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "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(
[
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys=["img", "seg"]),
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.5, spatial_axes=[0, 2]),
ToTensord(keys=["img", "seg"]),
]
)
val_transforms = Compose(
[
LoadNiftid(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
ScaleIntensityd(keys=["img", "seg"]),
ToTensord(keys=["img", "seg"]),
]
)
# define dataset, data loader
check_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
check_loader = DataLoader(
check_ds, batch_size=2, num_workers=4, collate_fn=list_data_collate, pin_memory=torch.cuda.is_available()
)
check_data = monai.utils.misc.first(check_loader)
print(check_data["img"].shape, check_data["seg"].shape)
# create a training data loader
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 = DataLoader(
train_ds,
batch_size=2,
shuffle=True,
num_workers=4,
collate_fn=list_data_collate,
pin_memory=torch.cuda.is_available(),
)
# create a validation data loader
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
val_loader = DataLoader(
val_ds, batch_size=5, num_workers=8, collate_fn=list_data_collate, pin_memory=torch.cuda.is_available()
)
# create UNet, DiceLoss and Adam optimizer
net = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
)
loss = monai.losses.DiceLoss(sigmoid=True)
lr = 1e-3
opt = torch.optim.Adam(net.parameters(), lr)
device = torch.device("cuda:0")
# Ignite trainer expects batch=(img, seg) and returns output=loss at every iteration,
# user can add output_transform to return other values, like: y_pred, y, etc.
def prepare_batch(batch, device=None, non_blocking=False):
return _prepare_batch((batch["img"], batch["seg"]), device, non_blocking)
trainer = create_supervised_trainer(net, opt, loss, device, False, prepare_batch=prepare_batch)
# adding checkpoint handler to save models (network params and optimizer stats) during training
checkpoint_handler = ModelCheckpoint("./runs/", "net", n_saved=10, require_empty=False)
trainer.add_event_handler(
event_name=Events.EPOCH_COMPLETED, handler=checkpoint_handler, to_save={"net": net, "opt": opt}
)
# StatsHandler prints loss at every iteration and print metrics at every epoch,
# we don't set metrics for trainer here, so just print loss, user can also customize print functions
# and can use output_transform to convert engine.state.output if it's not loss value
train_stats_handler = StatsHandler(name="trainer")
train_stats_handler.attach(trainer)
# TensorBoardStatsHandler plots loss at every iteration and plots metrics at every epoch, same as StatsHandler
train_tensorboard_stats_handler = TensorBoardStatsHandler()
train_tensorboard_stats_handler.attach(trainer)
validation_every_n_iters = 5
# set parameters for validation
metric_name = "Mean_Dice"
# add evaluation metric to the evaluator engine
val_metrics = {metric_name: MeanDice(sigmoid=True, to_onehot_y=False)}
# Ignite evaluator expects batch=(img, seg) and returns output=(y_pred, y) at every iteration,
# user can add output_transform to return other values
evaluator = create_supervised_evaluator(net, val_metrics, device, True, prepare_batch=prepare_batch)
@trainer.on(Events.ITERATION_COMPLETED(every=validation_every_n_iters))
def run_validation(engine):
evaluator.run(val_loader)
# add early stopping handler to evaluator
early_stopper = EarlyStopping(patience=4, score_function=stopping_fn_from_metric(metric_name), trainer=trainer)
evaluator.add_event_handler(event_name=Events.EPOCH_COMPLETED, handler=early_stopper)
# add stats event handler to print validation stats via evaluator
val_stats_handler = StatsHandler(
name="evaluator",
output_transform=lambda x: None, # no need to print loss value, so disable per iteration output
global_epoch_transform=lambda x: trainer.state.epoch,
) # fetch global epoch number from trainer
val_stats_handler.attach(evaluator)
# add handler to record metrics to TensorBoard at every validation epoch
val_tensorboard_stats_handler = TensorBoardStatsHandler(
output_transform=lambda x: None, # no need to plot loss value, so disable per iteration output
global_epoch_transform=lambda x: trainer.state.iteration,
) # fetch global iteration number from trainer
val_tensorboard_stats_handler.attach(evaluator)
# add handler to draw the first image and the corresponding label and model output in the last batch
# here we draw the 3D output as GIF format along the depth axis, every 2 validation iterations.
val_tensorboard_image_handler = TensorBoardImageHandler(
batch_transform=lambda batch: (batch["img"], batch["seg"]),
output_transform=lambda output: predict_segmentation(output[0]),
global_iter_transform=lambda x: trainer.state.epoch,
)
evaluator.add_event_handler(event_name=Events.ITERATION_COMPLETED(every=2), handler=val_tensorboard_image_handler)
train_epochs = 5
state = trainer.run(train_loader, train_epochs)
print(state)
shutil.rmtree(tempdir)
def test_invert(self):
set_determinism(seed=0)
im_fname, seg_fname = [
make_nifti_image(i)
for i in create_test_image_3d(101, 100, 107, noise_max=100)
]
transform = Compose([
LoadImaged(KEYS),
AddChanneld(KEYS),
Orientationd(KEYS, "RPS"),
Spacingd(KEYS,
pixdim=(1.2, 1.01, 0.9),
mode=["bilinear", "nearest"],
dtype=np.float32),
ScaleIntensityd("image", minv=1, maxv=10),
RandFlipd(KEYS, prob=0.5, spatial_axis=[1, 2]),
RandAxisFlipd(KEYS, prob=0.5),
RandRotate90d(KEYS, spatial_axes=(1, 2)),
RandZoomd(KEYS,
prob=0.5,
min_zoom=0.5,
max_zoom=1.1,
keep_size=True),
RandRotated(KEYS,
prob=0.5,
range_x=np.pi,
mode="bilinear",
align_corners=True),
RandAffined(KEYS, prob=0.5, rotate_range=np.pi, mode="nearest"),
ResizeWithPadOrCropd(KEYS, 100),
ToTensord(KEYS),
CastToTyped(KEYS, dtype=torch.uint8),
])
data = [{"image": im_fname, "label": seg_fname} for _ in range(12)]
# num workers = 0 for mac or gpu transforms
num_workers = 0 if sys.platform == "darwin" or torch.cuda.is_available(
) else 2
dataset = CacheDataset(data, transform=transform, progress=False)
loader = DataLoader(dataset, num_workers=num_workers, batch_size=5)
# set up engine
def _train_func(engine, batch):
self.assertTupleEqual(batch["image"].shape[1:], (1, 100, 100, 100))
engine.state.output = batch
engine.fire_event(IterationEvents.MODEL_COMPLETED)
return engine.state.output
engine = Engine(_train_func)
engine.register_events(*IterationEvents)
# set up testing handler
TransformInverter(
transform=transform,
loader=loader,
output_keys=["image", "label"],
batch_keys="label",
nearest_interp=True,
num_workers=0
if sys.platform == "darwin" or torch.cuda.is_available() else 2,
).attach(engine)
engine.run(loader, max_epochs=1)
set_determinism(seed=None)
self.assertTupleEqual(engine.state.output["image"].shape,
(2, 1, 100, 100, 100))
self.assertTupleEqual(engine.state.output["label"].shape,
(2, 1, 100, 100, 100))
for i in engine.state.output["image_inverted"] + engine.state.output[
"label_inverted"]:
torch.testing.assert_allclose(
i.to(torch.uint8).to(torch.float), i.to(torch.float))
self.assertTupleEqual(i.shape, (1, 100, 101, 107))
# check labels match
reverted = engine.state.output["label_inverted"][-1].detach().cpu(
).numpy()[0].astype(np.int32)
original = LoadImaged(KEYS)(data[-1])["label"]
n_good = np.sum(np.isclose(reverted, original, atol=1e-3))
reverted_name = engine.state.output["label_meta_dict"][
"filename_or_obj"][-1]
original_name = data[-1]["label"]
self.assertEqual(reverted_name, original_name)
print("invert diff", reverted.size - n_good)
self.assertTrue((reverted.size - n_good) in (25300, 1812),
"diff. in two possible values")