def test_compute(self, input_params, expected):
data = [
{
"image": torch.tensor([[[[2.0], [3.0]]]]),
"filename": "test1"
},
{
"image": torch.tensor([[[[6.0], [8.0]]]]),
"filename": "test2"
},
]
# set up engine, PostProcessing handler works together with post_transform of engine
engine = SupervisedEvaluator(
device=torch.device("cpu:0"),
val_data_loader=data,
epoch_length=2,
network=torch.nn.PReLU(),
post_transform=Compose([Activationsd(keys="pred", sigmoid=True)]),
val_handlers=[PostProcessing(**input_params)],
)
engine.run()
torch.testing.assert_allclose(engine.state.output["pred"], expected)
filename = engine.state.output.get("filename_bak")
if filename is not None:
self.assertEqual(filename, "test2")
def test_value_shape(self, input_param, test_input, output, expected_shape):
result = Activationsd(**input_param)(test_input)
assert_allclose(result["pred"], output["pred"], rtol=1e-3, type_test="tensor")
self.assertTupleEqual(result["pred"].shape, expected_shape)
if "label" in result:
assert_allclose(result["label"], output["label"], rtol=1e-3, type_test="tensor")
self.assertTupleEqual(result["label"].shape, expected_shape)
def run_interaction(self, train, compose):
data = [{
"image": np.ones((1, 2, 2, 2)).astype(np.float32),
"label": np.ones((1, 2, 2, 2))
} for _ in range(5)]
network = torch.nn.Linear(2, 2)
lr = 1e-3
opt = torch.optim.SGD(network.parameters(), lr)
loss = torch.nn.L1Loss()
train_transforms = Compose([
FindAllValidSlicesd(label="label", sids="sids"),
AddInitialSeedPointd(label="label",
guidance="guidance",
sids="sids"),
AddGuidanceSignald(image="image", guidance="guidance"),
ToTensord(keys=("image", "label")),
])
dataset = Dataset(data, transform=train_transforms)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=5)
iteration_transforms = [
Activationsd(keys="pred", sigmoid=True),
ToNumpyd(keys=["image", "label", "pred"]),
FindDiscrepancyRegionsd(label="label",
pred="pred",
discrepancy="discrepancy"),
AddRandomGuidanced(guidance="guidance",
discrepancy="discrepancy",
probability="probability"),
AddGuidanceSignald(image="image", guidance="guidance"),
ToTensord(keys=("image", "label")),
]
iteration_transforms = Compose(
iteration_transforms) if compose else iteration_transforms
i = Interaction(transforms=iteration_transforms,
train=train,
max_interactions=5)
self.assertEqual(len(i.transforms.transforms), 6,
"Mismatch in expected transforms")
# set up engine
engine = SupervisedTrainer(
device=torch.device("cpu"),
max_epochs=1,
train_data_loader=data_loader,
network=network,
optimizer=opt,
loss_function=loss,
iteration_update=i,
)
engine.add_event_handler(IterationEvents.INNER_ITERATION_STARTED,
add_one)
engine.add_event_handler(IterationEvents.INNER_ITERATION_COMPLETED,
add_one)
engine.run()
self.assertIsNotNone(engine.state.batch[0].get("guidance"),
"guidance is missing")
self.assertEqual(engine.state.best_metric, 9)
def test_value_shape(self, input_param, test_input, output, expected_shape):
result = Activationsd(**input_param)(test_input)
torch.testing.assert_allclose(result["pred"], output["pred"])
self.assertTupleEqual(result["pred"].shape, expected_shape)
if "label" in result:
torch.testing.assert_allclose(result["label"], output["label"])
self.assertTupleEqual(result["label"].shape, expected_shape)
def run_interaction(self, train, compose):
data = []
for i in range(5):
data.append({"image": torch.tensor([float(i)]), "label": torch.tensor([float(i)])})
network = torch.nn.Linear(1, 1)
lr = 1e-3
opt = torch.optim.SGD(network.parameters(), lr)
loss = torch.nn.L1Loss()
dataset = Dataset(data, transform=None)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=5)
iteration_transforms = [Activationsd(keys="pred", sigmoid=True), ToNumpyd(keys="pred")]
iteration_transforms = Compose(iteration_transforms) if compose else iteration_transforms
i = Interaction(transforms=iteration_transforms, train=train, max_interactions=5)
self.assertEqual(len(i.transforms.transforms), 2, "Mismatch in expected transforms")
# set up engine
engine = SupervisedTrainer(
device=torch.device("cpu"),
max_epochs=1,
train_data_loader=data_loader,
network=network,
optimizer=opt,
loss_function=loss,
iteration_update=i,
)
engine.add_event_handler(IterationEvents.INNER_ITERATION_STARTED, add_one)
engine.add_event_handler(IterationEvents.INNER_ITERATION_COMPLETED, add_one)
engine.run()
self.assertIsNotNone(engine.state.batch.get("probability"), "Probability is missing")
self.assertEqual(engine.state.best_metric, 9)
def test_compute(self, input_params, decollate, expected):
data = [
{"image": torch.tensor([[[[2.0], [3.0]]]]), "filename": ["test1"]},
{"image": torch.tensor([[[[6.0], [8.0]]]]), "filename": ["test2"]},
]
# set up engine, PostProcessing handler works together with postprocessing transforms of engine
engine = SupervisedEvaluator(
device=torch.device("cpu:0"),
val_data_loader=data,
epoch_length=2,
network=torch.nn.PReLU(),
postprocessing=Compose([Activationsd(keys="pred", sigmoid=True)]),
val_handlers=[PostProcessing(**input_params)],
decollate=decollate,
)
engine.run()
if isinstance(engine.state.output, list):
# test decollated list items
for o, e in zip(engine.state.output, expected):
torch.testing.assert_allclose(o["pred"], e)
filename = o.get("filename_bak")
if filename is not None:
self.assertEqual(filename, "test2")
else:
# test batch data
torch.testing.assert_allclose(engine.state.output["pred"], expected)
def post_transforms(self, data=None):
return [
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True, logit_thresh=0.5),
ToNumpyd(keys="pred"),
RestoreLabeld(keys="pred", ref_image="image", mode="nearest"),
AsChannelLastd(keys="pred"),
]
def train_post_transforms(self, context: Context):
return [
Activationsd(keys="pred", softmax=len(self.labels) > 1, sigmoid=len(self.labels) == 1),
AsDiscreted(
keys=("pred", "label"),
argmax=(True, False),
to_onehot=(len(self.labels) + 1, len(self.labels) + 1),
),
]
def post_transforms(self, data=None) -> Sequence[Callable]:
return [
EnsureTyped(keys="pred", device=data.get("device") if data else None),
Activationsd(keys="pred", softmax=True),
AsDiscreted(keys="pred", argmax=True),
SqueezeDimd(keys="pred", dim=0),
ToNumpyd(keys="pred"),
Restored(keys="pred", ref_image="image"),
]
def post_transforms(self, data=None) -> Sequence[Callable]:
return [
EnsureTyped(keys="pred",
device=data.get("device") if data else None),
Activationsd(keys="pred", softmax=True),
AsDiscreted(keys="pred", argmax=True),
ToNumpyd(keys="pred"),
Restored(keys="pred", ref_image="image"),
BoundingBoxd(keys="pred", result="result", bbox="bbox"),
]
def post_transforms(self, data=None) -> Sequence[Callable]:
return [
EnsureTyped(keys="pred",
device=data.get("device") if data else None),
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True, logit_thresh=0.5),
ToNumpyd(keys="pred"),
RestoreLabeld(keys="pred", ref_image="image", mode="nearest"),
AsChannelLastd(keys="pred"),
]
def post_transforms(self, data=None) -> Sequence[Callable]:
return [
EnsureTyped(keys="pred", device=data.get("device") if data else None),
Activationsd(keys="pred", softmax=len(self.labels) > 1, sigmoid=len(self.labels) == 1),
AsDiscreted(keys="pred", argmax=len(self.labels) > 1, threshold=0.5 if len(self.labels) == 1 else None),
SqueezeDimd(keys="pred", dim=0),
ToNumpyd(keys=("image", "pred")),
PostFilterLabeld(keys="pred", image="image"),
FindContoursd(keys="pred", labels=self.labels),
]
def train_post_transforms(self, context: Context):
return [
ToTensord(keys=("pred", "label")),
Activationsd(keys="pred", softmax=True),
AsDiscreted(
keys=("pred", "label"),
argmax=(True, False),
to_onehot=True,
n_classes=2,
),
]
def train_post_transforms(self, context: Context):
return [
Activationsd(keys="pred", softmax=True),
AsDiscreted(
keys=("pred", "label"),
argmax=(True, False),
to_onehot=(True, True),
n_classes=len(self._labels),
),
SplitPredsLabeld(keys="pred"),
]
def get_click_transforms(self, context: Context):
return [
Activationsd(keys="pred", sigmoid=True),
ToNumpyd(keys=("image", "label", "pred")),
FindDiscrepancyRegionsd(label="label",
pred="pred",
discrepancy="discrepancy"),
AddRandomGuidanced(guidance="guidance",
discrepancy="discrepancy",
probability="probability"),
AddGuidanceSignald(image="image", guidance="guidance"),
ToTensord(keys=("image", "label")),
]
def get_click_transforms():
return Compose([
Activationsd(keys="pred", sigmoid=True),
ToNumpyd(keys=("image", "label", "pred")),
FindDiscrepancyRegionsd(label="label",
pred="pred",
discrepancy="discrepancy"),
AddRandomGuidanced(
guidance="guidance",
discrepancy="discrepancy",
probability="probability",
),
AddGuidanceSignald(image="image", guidance="guidance"),
EnsureTyped(keys=("image", "label")),
])
def get_click_transforms():
return Compose([
Activationsd(keys='pred', sigmoid=True),
ToNumpyd(keys=('image', 'label', 'pred', 'probability', 'guidance')),
FindDiscrepancyRegionsd(label='label',
pred='pred',
discrepancy='discrepancy',
batched=True),
AddRandomGuidanced(guidance='guidance',
discrepancy='discrepancy',
probability='probability',
batched=True),
AddGuidanceSignald(image='image', guidance='guidance', batched=True),
ToTensord(keys=('image', 'label'))
])
def test_compute(self):
data = [
{
"image": torch.tensor([[[[2.0], [3.0]]]]),
"filename": ["test1"]
},
{
"image": torch.tensor([[[[6.0], [8.0]]]]),
"filename": ["test2"]
},
]
handlers = [
DecollateBatch(event="MODEL_COMPLETED"),
PostProcessing(transform=Compose([
Activationsd(keys="pred", sigmoid=True),
CopyItemsd(keys="filename", times=1, names="filename_bak"),
AsDiscreted(keys="pred",
threshold_values=True,
to_onehot=True,
num_classes=2),
])),
]
# set up engine, PostProcessing handler works together with postprocessing transforms of engine
engine = SupervisedEvaluator(
device=torch.device("cpu:0"),
val_data_loader=data,
epoch_length=2,
network=torch.nn.PReLU(),
# set decollate=False and execute some postprocessing first, then decollate in handlers
postprocessing=lambda x: dict(pred=x["pred"] + 1.0),
decollate=False,
val_handlers=handlers,
)
engine.run()
expected = torch.tensor([[[[1.0], [1.0]], [[0.0], [0.0]]]])
for o, e in zip(engine.state.output, expected):
torch.testing.assert_allclose(o["pred"], e)
filename = o.get("filename_bak")
if filename is not None:
self.assertEqual(filename, "test2")
def get_click_transforms(self, context: Context):
return [
Activationsd(keys="pred", softmax=True),
AsDiscreted(keys="pred", argmax=True),
ToNumpyd(keys=("image", "label", "pred")),
# Transforms for click simulation
FindDiscrepancyRegionsCustomd(keys="label",
pred="pred",
discrepancy="discrepancy"),
AddRandomGuidanceCustomd(
keys="NA",
guidance="guidance",
discrepancy="discrepancy",
probability="probability",
),
AddGuidanceSignalCustomd(
keys="image",
guidance="guidance",
number_intensity_ch=self.number_intensity_ch),
#
ToTensord(keys=("image", "label")),
]
def post_transforms(self, data=None) -> Sequence[Callable]:
largest_cc = False if not data else data.get("largest_cc", False)
applied_labels = list(self.labels.values()) if isinstance(
self.labels, dict) else self.labels
t = [
EnsureTyped(keys="pred",
device=data.get("device") if data else None),
Activationsd(keys="pred",
softmax=len(self.labels) > 1,
sigmoid=len(self.labels) == 1),
AsDiscreted(keys="pred",
argmax=len(self.labels) > 1,
threshold=0.5 if len(self.labels) == 1 else None),
]
if largest_cc:
t.append(
KeepLargestConnectedComponentd(keys="pred",
applied_labels=applied_labels))
t.extend([
ToNumpyd(keys="pred"),
Restored(keys="pred", ref_image="image"),
])
return t
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(
[
LoadNiftid(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 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 = [{"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:
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
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.config.get_torch_version_tuple() >= (1, 6) else False,
)
evaluator.run()
dist.destroy_process_group()
def main(tempdir):
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
################################ DATASET ################################
# 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, 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(
[
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]),
ToTensord(keys=["image", "label"]),
]
)
val_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys="image"),
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)
################################ DATASET ################################
################################ NETWORK ################################
# 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)
################################ NETWORK ################################
################################ LOSS ################################
loss = monai.losses.DiceLoss(sigmoid=True)
################################ LOSS ################################
################################ OPT ################################
opt = torch.optim.Adam(net.parameters(), 1e-3)
################################ OPT ################################
################################ LR ################################
lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1)
################################ LR ################################
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.utils.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.utils.get_torch_version_tuple() >= (1, 6) else False,
)
trainer.run()
def test_none_postfix(self, input_param, test_input, output,
expected_shape):
result = Activationsd(**input_param)(test_input)
torch.testing.assert_allclose(result["pred"], output["pred"])
self.assertTupleEqual(result["pred"].shape, expected_shape)
def train(self,
train_info,
valid_info,
hyperparameters,
run_data_check=False):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
if not run_data_check:
start_dt = datetime.datetime.now()
start_dt_string = start_dt.strftime('%d/%m/%Y %H:%M:%S')
print(f'Training started: {start_dt_string}')
# 1. Create folders to save the model
timedate_info = str(
datetime.datetime.now()).split(' ')[0] + '_' + str(
datetime.datetime.now().strftime("%H:%M:%S")).replace(
':', '-')
path_to_model = os.path.join(
self.out_dir, 'trained_models',
self.unique_name + '_' + timedate_info)
os.mkdir(path_to_model)
# 2. Load hyperparameters
learning_rate = hyperparameters['learning_rate']
weight_decay = hyperparameters['weight_decay']
total_epoch = hyperparameters['total_epoch']
multiplicator = hyperparameters['multiplicator']
batch_size = hyperparameters['batch_size']
validation_epoch = hyperparameters['validation_epoch']
validation_interval = hyperparameters['validation_interval']
H = hyperparameters['H']
L = hyperparameters['L']
# 3. Consider class imbalance
negative, positive = 0, 0
for _, label in train_info:
if int(label) == 0:
negative += 1
elif int(label) == 1:
positive += 1
pos_weight = torch.Tensor([(negative / positive)]).to(self.device)
# 4. Create train and validation loaders, batch_size = 10 for validation loader (10 central slices)
train_data = get_data_from_info(self.image_data_dir, self.seg_data_dir,
train_info)
valid_data = get_data_from_info(self.image_data_dir, self.seg_data_dir,
valid_info)
large_image_splitter(train_data, self.cache_dir)
set_determinism(seed=100)
train_trans, valid_trans = self.transformations(H, L)
train_dataset = PersistentDataset(
data=train_data[:],
transform=train_trans,
cache_dir=self.persistent_dataset_dir)
valid_dataset = PersistentDataset(
data=valid_data[:],
transform=valid_trans,
cache_dir=self.persistent_dataset_dir)
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=self.pin_memory,
num_workers=self.num_workers,
collate_fn=PadListDataCollate(
Method.SYMMETRIC, NumpyPadMode.CONSTANT))
valid_loader = DataLoader(valid_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=self.pin_memory,
num_workers=self.num_workers,
collate_fn=PadListDataCollate(
Method.SYMMETRIC, NumpyPadMode.CONSTANT))
# Perform data checks
if run_data_check:
check_data = monai.utils.misc.first(train_loader)
print(check_data["image"].shape, check_data["label"])
for i in range(batch_size):
multi_slice_viewer(
check_data["image"][i, 0, :, :, :],
check_data["image_meta_dict"]["filename_or_obj"][i])
exit()
"""c = 1
for d in train_loader:
img = d["image"]
seg = d["seg"][0]
seg, _ = nrrd.read(seg)
img_name = d["image_meta_dict"]["filename_or_obj"][0]
print(c, "Name:", img_name, "Size:", img.nelement()*img.element_size()/1024/1024, "MB", "shape:", img.shape)
multi_slice_viewer(img[0, 0, :, :, :], d["image_meta_dict"]["filename_or_obj"][0])
#multi_slice_viewer(seg, d["image_meta_dict"]["filename_or_obj"][0])
c += 1
exit()"""
# 5. Prepare model
model = ModelCT().to(self.device)
# 6. Define loss function, optimizer and scheduler
loss_function = torch.nn.BCEWithLogitsLoss(
pos_weight) # pos_weight for class imbalance
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer,
multiplicator,
last_epoch=-1)
# 7. Create post validation transforms and handlers
path_to_tensorboard = os.path.join(self.out_dir, 'tensorboard')
writer = SummaryWriter(log_dir=path_to_tensorboard)
valid_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
])
valid_handlers = [
StatsHandler(output_transform=lambda x: None),
TensorBoardStatsHandler(summary_writer=writer,
output_transform=lambda x: None),
CheckpointSaver(save_dir=path_to_model,
save_dict={"model": model},
save_key_metric=True),
MetricsSaver(save_dir=path_to_model,
metrics=['Valid_AUC', 'Valid_ACC']),
]
# 8. Create validatior
discrete = AsDiscrete(threshold_values=True)
evaluator = SupervisedEvaluator(
device=self.device,
val_data_loader=valid_loader,
network=model,
post_transform=valid_post_transforms,
key_val_metric={
"Valid_AUC":
ROCAUC(output_transform=lambda x: (x["pred"], x["label"]))
},
additional_metrics={
"Valid_Accuracy":
Accuracy(output_transform=lambda x:
(discrete(x["pred"]), x["label"]))
},
val_handlers=valid_handlers,
amp=self.amp,
)
# 9. Create trainer
# Loss function does the last sigmoid, so we dont need it here.
train_post_transforms = Compose([
# Empty
])
logger = MetricLogger(evaluator=evaluator)
train_handlers = [
logger,
LrScheduleHandler(lr_scheduler=scheduler, print_lr=True),
ValidationHandlerCT(validator=evaluator,
start=validation_epoch,
interval=validation_interval,
epoch_level=True),
StatsHandler(tag_name="loss",
output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(summary_writer=writer,
tag_name="Train_Loss",
output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir=path_to_model,
save_dict={
"model": model,
"opt": optimizer
},
save_interval=1,
n_saved=1),
]
trainer = SupervisedTrainer(
device=self.device,
max_epochs=total_epoch,
train_data_loader=train_loader,
network=model,
optimizer=optimizer,
loss_function=loss_function,
post_transform=train_post_transforms,
train_handlers=train_handlers,
amp=self.amp,
)
# 10. Run trainer
trainer.run()
# 11. Save results
np.save(path_to_model + '/AUCS.npy',
np.array(logger.metrics['Valid_AUC']))
np.save(path_to_model + '/ACCS.npy',
np.array(logger.metrics['Valid_ACC']))
np.save(path_to_model + '/LOSSES.npy', np.array(logger.loss))
np.save(path_to_model + '/PARAMETERS.npy', np.array(hyperparameters))
return path_to_model
def run_training_test(root_dir, device="cuda:0", amp=False):
images = sorted(glob(os.path.join(root_dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(root_dir, "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
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=root_dir,
output_transform=lambda x: None),
TensorBoardImageHandler(log_dir=root_dir,
batch_transform=lambda x:
(x["image"], x["label"]),
output_transform=lambda x: x["pred"]),
CheckpointSaver(save_dir=root_dir,
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,
amp=True if amp 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=root_dir,
tag_name="train_loss",
output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir=root_dir,
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,
amp=True if amp else False,
)
trainer.run()
return evaluator.state.best_metric
def run_inference_test(root_dir, model_file, device="cuda:0", amp=False):
images = sorted(glob(os.path.join(root_dir, "im*.nii.gz")))
segs = sorted(glob(os.path.join(root_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([
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
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=f"{model_file}", load_dict={"net": net}),
SegmentationSaver(
output_dir=root_dir,
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,
amp=True if amp else False,
)
evaluator.run()
return evaluator.state.best_metric
def main(train_output):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print_config()
# Setup directories
dirs = setup_directories()
# Setup torch device
device, using_gpu = create_device("cuda")
# Load and randomize images
# HACKATON image and segmentation data
hackathon_dir = os.path.join(dirs["data"], 'HACKATHON')
map_fn = lambda x: (x[0], int(x[1]))
with open(os.path.join(hackathon_dir, "train.txt"), 'r') as fp:
train_info_hackathon = [
map_fn(entry.strip().split(',')) for entry in fp.readlines()
]
image_dir = os.path.join(hackathon_dir, 'images', 'train')
seg_dir = os.path.join(hackathon_dir, 'segmentations', 'train')
_train_data_hackathon = get_data_from_info(image_dir,
seg_dir,
train_info_hackathon,
dual_output=False)
large_image_splitter(_train_data_hackathon, dirs["cache"])
balance_training_data(_train_data_hackathon, seed=72)
# PSUF data
"""psuf_dir = os.path.join(dirs["data"], 'psuf')
with open(os.path.join(psuf_dir, "train.txt"), 'r') as fp:
train_info = [entry.strip().split(',') for entry in fp.readlines()]
image_dir = os.path.join(psuf_dir, 'images')
train_data_psuf = get_data_from_info(image_dir, None, train_info)"""
# Split data into train, validate and test
train_split, test_data_hackathon = train_test_split(_train_data_hackathon,
test_size=0.2,
shuffle=True,
random_state=42)
#train_data_hackathon, valid_data_hackathon = train_test_split(train_split, test_size=0.2, shuffle=True, random_state=43)
# Setup transforms
# Crop foreground
crop_foreground = CropForegroundd(
keys=["image"],
source_key="image",
margin=(5, 5, 0),
#select_fn = lambda x: x != 0
)
# Crop Z
crop_z = RelativeCropZd(keys=["image"], relative_z_roi=(0.07, 0.12))
# Window width and level (window center)
WW, WL = 1500, -600
ct_window = CTWindowd(keys=["image"], width=WW, level=WL)
spatial_pad = SpatialPadd(keys=["image"], spatial_size=(-1, -1, 30))
resize = Resized(keys=["image"],
spatial_size=(int(512 * 0.50), int(512 * 0.50), -1),
mode="trilinear")
# Create transforms
common_transform = Compose([
LoadImaged(keys=["image"]),
ct_window,
CTSegmentation(keys=["image"]),
AddChanneld(keys=["image"]),
resize,
crop_foreground,
crop_z,
spatial_pad,
])
hackathon_train_transfrom = Compose([
common_transform,
ToTensord(keys=["image"]),
]).flatten()
psuf_transforms = Compose([
LoadImaged(keys=["image"]),
AddChanneld(keys=["image"]),
ToTensord(keys=["image"]),
])
# Setup data
#set_determinism(seed=100)
test_dataset = PersistentDataset(data=test_data_hackathon[:],
transform=hackathon_train_transfrom,
cache_dir=dirs["persistent"])
test_loader = DataLoader(test_dataset,
batch_size=2,
shuffle=True,
pin_memory=using_gpu,
num_workers=1,
collate_fn=PadListDataCollate(
Method.SYMMETRIC, NumpyPadMode.CONSTANT))
# Setup network, loss function, optimizer and scheduler
network = nets.DenseNet121(spatial_dims=3, in_channels=1,
out_channels=1).to(device)
# Setup validator and trainer
valid_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
])
# Setup tester
tester = Tester(device=device,
test_data_loader=test_loader,
load_dir=train_output,
out_dir=dirs["out"],
network=network,
post_transform=valid_post_transforms,
non_blocking=using_gpu,
amp=using_gpu)
# Run tester
tester.run()
def configure(self):
self.set_device()
network = UNet(
dimensions=3,
in_channels=1,
out_channels=2,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
).to(self.device)
if self.multi_gpu:
network = DistributedDataParallel(
module=network,
device_ids=[self.device],
find_unused_parameters=False,
)
train_transforms = Compose([
LoadImaged(keys=("image", "label")),
EnsureChannelFirstd(keys=("image", "label")),
Spacingd(keys=("image", "label"),
pixdim=[1.0, 1.0, 1.0],
mode=["bilinear", "nearest"]),
ScaleIntensityRanged(
keys="image",
a_min=-57,
a_max=164,
b_min=0.0,
b_max=1.0,
clip=True,
),
CropForegroundd(keys=("image", "label"), source_key="image"),
RandCropByPosNegLabeld(
keys=("image", "label"),
label_key="label",
spatial_size=(96, 96, 96),
pos=1,
neg=1,
num_samples=4,
image_key="image",
image_threshold=0,
),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
ToTensord(keys=("image", "label")),
])
train_datalist = load_decathlon_datalist(self.data_list_file_path,
True, "training")
if self.multi_gpu:
train_datalist = partition_dataset(
data=train_datalist,
shuffle=True,
num_partitions=dist.get_world_size(),
even_divisible=True,
)[dist.get_rank()]
train_ds = CacheDataset(
data=train_datalist,
transform=train_transforms,
cache_num=32,
cache_rate=1.0,
num_workers=4,
)
train_data_loader = DataLoader(
train_ds,
batch_size=2,
shuffle=True,
num_workers=4,
)
val_transforms = Compose([
LoadImaged(keys=("image", "label")),
EnsureChannelFirstd(keys=("image", "label")),
ScaleIntensityRanged(
keys="image",
a_min=-57,
a_max=164,
b_min=0.0,
b_max=1.0,
clip=True,
),
CropForegroundd(keys=("image", "label"), source_key="image"),
ToTensord(keys=("image", "label")),
])
val_datalist = load_decathlon_datalist(self.data_list_file_path, True,
"validation")
val_ds = CacheDataset(val_datalist, val_transforms, 9, 0.0, 4)
val_data_loader = DataLoader(
val_ds,
batch_size=1,
shuffle=False,
num_workers=4,
)
post_transform = Compose([
Activationsd(keys="pred", softmax=True),
AsDiscreted(
keys=["pred", "label"],
argmax=[True, False],
to_onehot=True,
n_classes=2,
),
])
# metric
key_val_metric = {
"val_mean_dice":
MeanDice(
include_background=False,
output_transform=lambda x: (x["pred"], x["label"]),
device=self.device,
)
}
val_handlers = [
StatsHandler(output_transform=lambda x: None),
CheckpointSaver(
save_dir=self.ckpt_dir,
save_dict={"model": network},
save_key_metric=True,
),
TensorBoardStatsHandler(log_dir=self.ckpt_dir,
output_transform=lambda x: None),
]
self.eval_engine = SupervisedEvaluator(
device=self.device,
val_data_loader=val_data_loader,
network=network,
inferer=SlidingWindowInferer(
roi_size=[160, 160, 160],
sw_batch_size=4,
overlap=0.5,
),
post_transform=post_transform,
key_val_metric=key_val_metric,
val_handlers=val_handlers,
amp=self.amp,
)
optimizer = torch.optim.Adam(network.parameters(), self.learning_rate)
loss_function = DiceLoss(to_onehot_y=True, softmax=True)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=5000,
gamma=0.1)
train_handlers = [
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
ValidationHandler(validator=self.eval_engine,
interval=self.val_interval,
epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(
log_dir=self.ckpt_dir,
tag_name="train_loss",
output_transform=lambda x: x["loss"],
),
]
self.train_engine = SupervisedTrainer(
device=self.device,
max_epochs=self.max_epochs,
train_data_loader=train_data_loader,
network=network,
optimizer=optimizer,
loss_function=loss_function,
inferer=SimpleInferer(),
post_transform=post_transform,
key_train_metric=None,
train_handlers=train_handlers,
amp=self.amp,
)
if self.local_rank > 0:
self.train_engine.logger.setLevel(logging.WARNING)
self.eval_engine.logger.setLevel(logging.WARNING)
def train(args):
if args.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"))
# 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 = [{
"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]),
ToTensord(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:{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)
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([
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),
]
if dist.get_rank() == 0:
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
train_handlers.extend([
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
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.config.get_torch_version_tuple() >=
(1, 6) else False,
post_transform=train_post_transforms,
key_train_metric={
"train_acc":
Accuracy(output_transform=lambda x: (x["pred"], x["label"]),
device=device)
},
train_handlers=train_handlers,
)
trainer.run()
dist.destroy_process_group()
