def prepare_data(self):
data_dir = self.hparams.data_dir
# Train imgs/masks
train_imgs = []
train_masks = []
with open(data_dir + 'train_imgs.txt', 'r') as f:
train_imgs = [data_dir + image.rstrip() for image in f.readlines()]
with open(data_dir + 'train_masks.txt', 'r') as f:
train_masks = [data_dir + mask.rstrip() for mask in f.readlines()]
train_dicts = [{'image': image, 'mask': mask} for (image, mask) in zip(train_imgs, train_masks)]
train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.2)
# Basic transforms
data_keys = ["image", "mask"]
data_transforms = Compose(
[
LoadNiftid(keys=data_keys),
AddChanneld(keys=data_keys),
ScaleIntensityRangePercentilesd(
keys='image',
lower=25,
upper=75,
b_min=-0.5,
b_max=0.5
)
]
)
self.train_dataset = monai.data.CacheDataset(
data=train_dicts,
transform=Compose(
[
data_transforms,
RandCropByPosNegLabeld(
keys=data_keys,
label_key="mask",
spatial_size=self.hparams.patch_size,
num_samples=4,
image_key="image",
pos=0.8,
neg=0.2
),
ToTensord(keys=data_keys)
]
),
cache_rate=1.0
)
self.val_dataset = monai.data.CacheDataset(
data=val_dicts,
transform=Compose(
[
data_transforms,
CenterSpatialCropd(keys=data_keys, roi_size=self.hparams.patch_size),
ToTensord(keys=data_keys)
]
),
cache_rate=1.0
)
def train_pre_transforms(self, context: Context):
return [
LoadImaged(keys=("image", "label")),
AddChanneld(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"),
EnsureTyped(keys=("image", "label"), device=context.device),
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),
SelectItemsd(keys=("image", "label")),
]
def test_type_shape(self, input_param, input_data, expected_type,
expected_shape):
result = RandCropByPosNegLabeld(**input_param)(input_data)
self.assertIsInstance(result, expected_type)
self.assertTupleEqual(result[0]["image"].shape, expected_shape)
self.assertTupleEqual(result[0]["extral"].shape, expected_shape)
self.assertTupleEqual(result[0]["label"].shape, expected_shape)
def test_type_shape(self, input_param, input_data, expected_type, expected_shape):
result = RandCropByPosNegLabeld(**input_param)(input_data)
self.assertIsInstance(result, expected_type)
self.assertTupleEqual(result[0]["image"].shape, expected_shape)
self.assertTupleEqual(result[0]["extral"].shape, expected_shape)
self.assertTupleEqual(result[0]["label"].shape, expected_shape)
for i, item in enumerate(result):
self.assertEqual(item["image_meta_dict"]["patch_index"], i)
def test_type_shape(self, input_param, input_data, expected_shape):
for p in TEST_NDARRAYS:
input_param_mod = self.convert_data_type(p, input_param)
input_data_mod = self.convert_data_type(p, input_data)
cropper = RandCropByPosNegLabeld(**input_param_mod)
cropper.set_random_state(0)
result = cropper(input_data_mod)
self.assertIsInstance(result, list)
_len = len(tuple(input_data.keys()))
self.assertTupleEqual(
tuple(result[0].keys())[:_len], tuple(input_data.keys()))
for k in ("image", "extra", "label"):
self.assertTupleEqual(result[0][k].shape, expected_shape)
for i, item in enumerate(result):
self.assertEqual(item[k + "_meta_dict"]["patch_index"], i)
def setup(self, stage):
data_dir = 'data/'
# Train imgs/masks
train_imgs = []
with open(data_dir + 'train_imgs.txt', 'r') as f:
train_imgs = [image.rstrip() for image in f.readlines()]
train_masks = []
with open(data_dir + 'train_masks.txt', 'r') as f:
train_masks = [mask.rstrip() for mask in f.readlines()]
train_dicts = [{'image': image, 'mask': mask} for (image, mask) in zip(train_imgs, train_masks)]
train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.2)
# Basic transforms
data_keys = ["image", "mask"]
data_transforms = Compose(
[
LoadNiftid(keys=data_keys),
AddChanneld(keys=data_keys),
NormalizeIntensityd(keys="image"),
RandCropByPosNegLabeld(
keys=data_keys, label_key="mask", size=(256, 256, 16), num_samples=4, image_key="image"
),
]
)
self.train_dataset = monai.data.CacheDataset(
data=train_dicts,
transform=Compose(
[
data_transforms,
self.augmentations,
ToTensord(keys=data_keys)
]
),
cache_rate=1.0
)
self.val_dataset = monai.data.CacheDataset(
data=val_dicts,
transform=Compose(
[
data_transforms,
ToTensord(keys=data_keys)
]
),
cache_rate=1.0
)
def get_xforms(args, mode="train", keys=("image", "label")):
"""returns a composed transform for train/val/infer."""
xforms = [
LoadNiftid(keys),
AddChanneld(keys),
Orientationd(keys, axcodes="LPS"),
Spacingd(keys,
pixdim=(1.25, 1.25, 5.0),
mode=("bilinear", "nearest")[:len(keys)]),
ScaleIntensityRanged(keys[0],
a_min=-1000.0,
a_max=500.0,
b_min=0.0,
b_max=1.0,
clip=True),
]
if mode == "train":
xforms.extend([
SpatialPadd(keys,
spatial_size=(args.patch_size, args.patch_size, -1),
mode="reflect"), # ensure at least 192x192
RandAffined(
keys,
prob=0.15,
rotate_range=(-0.05, 0.05),
scale_range=(-0.1, 0.1),
mode=("bilinear", "nearest"),
as_tensor_output=False,
),
RandCropByPosNegLabeld(keys,
label_key=keys[1],
spatial_size=(args.patch_size,
args.patch_size,
args.n_slice),
num_samples=3),
RandGaussianNoised(keys[0], prob=0.15, std=0.01),
RandFlipd(keys, spatial_axis=0, prob=0.5),
RandFlipd(keys, spatial_axis=1, prob=0.5),
RandFlipd(keys, spatial_axis=2, prob=0.5),
])
dtype = (np.float32, np.uint8)
if mode == "val":
dtype = (np.float32, np.uint8)
if mode == "infer":
dtype = (np.float32, )
xforms.extend([CastToTyped(keys, dtype=dtype), ToTensord(keys)])
return monai.transforms.Compose(xforms)
def prepare_data(self):
data_df = pd.read_csv(
'/data/shared/prostate/yale_prostate/input_lists/MR_yale.csv')
train_imgs = data_df['IMAGE'][0:295].tolist()
train_masks = data_df['SEGM'][0:295].tolist()
train_dicts = [{
'image': image,
'mask': mask
} for (image, mask) in zip(train_imgs, train_masks)]
train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.2)
# Basic transforms
data_keys = ["image", "mask"]
data_transforms = Compose([
LoadNiftid(keys=data_keys),
AddChanneld(keys=data_keys),
ScaleIntensityRangePercentilesd(keys='image',
lower=25,
upper=75,
b_min=-0.5,
b_max=0.5)
])
self.train_dataset = monai.data.CacheDataset(
data=train_dicts,
transform=Compose([
data_transforms,
RandCropByPosNegLabeld(keys=data_keys,
label_key="mask",
spatial_size=self.hparams.patch_size,
num_samples=4,
image_key="image",
pos=0.8,
neg=0.2),
ToTensord(keys=data_keys)
]),
cache_rate=1.0)
self.val_dataset = monai.data.CacheDataset(
data=val_dicts,
transform=Compose([
data_transforms,
CenterSpatialCropd(keys=data_keys,
roi_size=self.hparams.patch_size),
ToTensord(keys=data_keys)
]),
cache_rate=1.0)
def setup(self, stage):
data_dir = "data/"
# Train imgs/masks
train_imgs = []
with open(data_dir + "train_imgs.txt", "r") as f:
train_imgs = [image.rstrip() for image in f.readlines()]
train_masks = []
with open(data_dir + "train_masks.txt", "r") as f:
train_masks = [mask.rstrip() for mask in f.readlines()]
train_dicts = [{
"image": image,
"mask": mask
} for (image, mask) in zip(train_imgs, train_masks)]
train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.2)
# Basic transforms
data_keys = ["image", "mask"]
data_transforms = Compose([
AddChanneld(keys=data_keys),
NormalizeIntensityd(keys="image"),
RandCropByPosNegLabeld(
keys=data_keys,
label_key="mask",
spatial_size=self.hparams.patch_size,
num_samples=4,
image_key="image",
),
])
self.train_dataset = monai.data.CacheDataset(
data=train_dicts,
transform=Compose([data_transforms,
ToTensord(keys=data_keys)]),
cache_rate=1.0,
)
self.val_dataset = monai.data.CacheDataset(
data=val_dicts,
transform=Compose([data_transforms,
ToTensord(keys=data_keys)]),
cache_rate=1.0,
)
def setup(self, stage):
data_df = pd.read_csv(
'/data/shared/prostate/yale_prostate/input_lists/MR_yale.csv')
train_imgs = data_df['IMAGE'][0:295].tolist()
train_masks = data_df['SEGM'][0:295].tolist()
train_dicts = [{
'image': image,
'mask': mask
} for (image, mask) in zip(train_imgs, train_masks)]
train_dicts, val_dicts = train_test_split(train_dicts, test_size=0.15)
# Basic transforms
data_keys = ["image", "mask"]
data_transforms = Compose([
LoadNiftid(keys=data_keys),
AddChanneld(keys=data_keys),
NormalizeIntensityd(keys="image"),
RandCropByPosNegLabeld(keys=data_keys,
label_key="mask",
spatial_size=self.hparams.patch_size,
num_samples=4,
image_key="image",
pos=0.7,
neg=0.3),
])
self.train_dataset = monai.data.CacheDataset(
data=train_dicts,
transform=Compose([data_transforms,
ToTensord(keys=data_keys)]),
cache_rate=1.0)
self.val_dataset = monai.data.CacheDataset(
data=val_dicts,
transform=Compose([data_transforms,
ToTensord(keys=data_keys)]),
cache_rate=1.0)
def get_data_loaders(batch_size, patch_size):
### Data collection
data_dir = "../../data/"
print("Available directories: ", os.listdir(data_dir))
# Get paths for images and masks, organize into dictionaries
images = sorted(glob.glob(data_dir + "**/*CTImg*", recursive=True))
masks = sorted(glob.glob(data_dir + "**/*Mask*", recursive=True))
# Dataset selection
train_dicts = select_animals(images, masks, [12, 13, 14, 18, 20])
val_dicts = select_animals(images, masks, [25])
data_keys = ["image", "mask"]
# Data transformation
data_transforms = Compose([
LoadNiftid(keys=data_keys),
AddChanneld(keys=data_keys),
NormalizeIntensityd(keys="image"),
RandCropByPosNegLabeld(keys=data_keys,
label_key="mask",
size=patch_size,
num_samples=16,
image_key="image"),
ToTensord(keys=data_keys),
])
# Data loaders
data_loaders = {
"train":
create_loader(
train_dicts,
batch_size=batch_size,
transforms=data_transforms,
shuffle=True,
),
"val":
create_loader(val_dicts, transforms=data_transforms),
}
for key in data_loaders:
print(key, len(data_loaders[key]))
return data_loaders
def get_xforms_with_synthesis(mode="synthesis", keys=("image", "label"), keys2=("image", "label", "synthetic_lesion")):
"""returns a composed transform for train/val/infer."""
xforms = [
LoadImaged(keys),
AddChanneld(keys),
Orientationd(keys, axcodes="LPS"),
Spacingd(keys, pixdim=(1.25, 1.25, 5.0), mode=("bilinear", "nearest")[: len(keys)]),
ScaleIntensityRanged(keys[0], a_min=-1000.0, a_max=500.0, b_min=0.0, b_max=1.0, clip=True),
CopyItemsd(keys,1, names=['image_1', 'label_1']),
]
if mode == "synthesis":
xforms.extend([
SpatialPadd(keys, spatial_size=(192, 192, -1), mode="reflect"), # ensure at least 192x192
RandCropByPosNegLabeld(keys, label_key=keys[1], spatial_size=(192, 192, 16), num_samples=3),
TransCustom(keys, path_synthesis, read_cea_aug_slice2,
pseudo_healthy_with_texture, scans_syns, decreasing_sequence, GEN=15,
POST_PROCESS=True, mask_outer_ring=True, new_value=.5),
RandAffined(
# keys,
keys2,
prob=0.15,
rotate_range=(0.05, 0.05, None), # 3 parameters control the transform on 3 dimensions
scale_range=(0.1, 0.1, None),
mode=("bilinear", "nearest", "bilinear"),
# mode=("bilinear", "nearest"),
as_tensor_output=False
),
RandGaussianNoised((keys2[0],keys2[2]), prob=0.15, std=0.01),
# RandGaussianNoised(keys[0], prob=0.15, std=0.01),
RandFlipd(keys, spatial_axis=0, prob=0.5),
RandFlipd(keys, spatial_axis=1, prob=0.5),
RandFlipd(keys, spatial_axis=2, prob=0.5),
TransCustom2(0.333)
])
dtype = (np.float32, np.uint8)
# dtype = (np.float32, np.uint8, np.float32)
xforms.extend([CastToTyped(keys, dtype=dtype)])
return monai.transforms.Compose(xforms)
def train_pre_transforms(self, context: Context):
return [
LoadImaged(keys=("image", "label"), reader="ITKReader"),
NormalizeLabelsInDatasetd(
keys="label",
label_names=self._labels), # Specially for missing labels
EnsureChannelFirstd(keys=("image", "label")),
Spacingd(keys=("image", "label"),
pixdim=self.target_spacing,
mode=("bilinear", "nearest")),
CropForegroundd(keys=("image", "label"), source_key="image"),
SpatialPadd(keys=("image", "label"),
spatial_size=self.spatial_size),
ScaleIntensityRanged(keys="image",
a_min=-175,
a_max=250,
b_min=0.0,
b_max=1.0,
clip=True),
RandCropByPosNegLabeld(
keys=("image", "label"),
label_key="label",
spatial_size=self.spatial_size,
pos=1,
neg=1,
num_samples=self.num_samples,
image_key="image",
image_threshold=0,
),
EnsureTyped(keys=("image", "label"), device=context.device),
RandFlipd(keys=("image", "label"), spatial_axis=[0], prob=0.10),
RandFlipd(keys=("image", "label"), spatial_axis=[1], prob=0.10),
RandFlipd(keys=("image", "label"), spatial_axis=[2], prob=0.10),
RandRotate90d(keys=("image", "label"), prob=0.10, max_k=3),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
SelectItemsd(keys=("image", "label")),
]
def _get_loader(self, folders):
images = []
segs = []
for folder in folders:
images += glob(os.path.join(folder, "*_im.nii.gz"))
segs += glob(os.path.join(folder, "*_seg.nii.gz"))
images = sorted(images, key=os.path.basename)
segs = sorted(segs, key=os.path.basename)
files = [{"img": img, "seg": seg} for img, seg in zip(images, segs)]
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]),
ToTensord(keys=["img", "seg"]),
])
ds = CacheDataset(data=files, transform=transforms)
loader = DataLoader(
ds,
batch_size=2,
shuffle=True,
num_workers=4,
collate_fn=list_data_collate,
pin_memory=torch.cuda.is_available(),
)
return loader
mode=("bilinear", "nearest")),
Orientationd(keys=["image", "label"], axcodes="RAS"),
ScaleIntensityRanged(
keys=["image"],
a_min=-300,
a_max=300,
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,
),
Rand3DElasticd(
keys=["image", "label"],
sigma_range=(0, 1),
magnitude_range=(0, 1),
spatial_size=None,
prob=0.5,
rotate_range=(0, -math.pi / 36, math.pi / 36, 0), # -15, 15 / -5, 5
shear_range=None,
translate_range=None,
scale_range=None,
mode=("bilinear", "nearest"),
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_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),
val_handlers=val_handlers,
key_val_metric={
"val_mean_dice":
MeanDice(include_background=True,
add_sigmoid=True,
output_transform=lambda x: (x[Keys.PRED], x[Keys.LABEL]))
},
additional_metrics=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(),
train_handlers=train_handlers,
amp=False,
key_train_metric=None,
)
trainer.run()
shutil.rmtree(tempdir)
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 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", 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=1, num_workers=4, collate_fn=list_data_collate, pin_memory=torch.cuda.is_available()
)
# 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(do_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["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}")
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["img"].to(device), val_data["seg"].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 = compute_meandice(
y_pred=val_outputs, y=val_labels, include_background=True, to_onehot_y=False, add_sigmoid=True
)
metric_count += len(value)
metric_sum += value.sum().item()
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 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 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"]),
])
# partition dataset based on current rank number, every rank trains with its own data
# it can avoid duplicated caching content in each rank, but will not do global shuffle before every epoch
data_part = partition_dataset(
data=train_files,
num_partitions=dist.get_world_size(),
shuffle=True,
even_divisible=True,
)[dist.get_rank()]
train_ds = SmartCacheDataset(
data=data_part,
transform=train_transforms,
replace_rate=0.2,
cache_num=
15, # we suppose to use 2 ranks in this example, every rank has 20 training images
num_init_workers=2,
num_replace_workers=2,
)
# 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=2,
pin_memory=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(
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()
# start the replacement thread of SmartCache
train_ds.start()
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["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 = math.ceil(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)
# replace 20% of cache content for next epoch
train_ds.update_cache()
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
# stop replacement thread of SmartCache
train_ds.shutdown()
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()
)
from DataSet import *
from VNet import *
from Training import *
train_transform = Compose([
LoadImaged(keys=[DataType.Image, DataType.Label]),
AddChanneld(keys=[DataType.Image, DataType.Label]),
CropForegroundd(keys=[DataType.Image, DataType.Label],
source_key=DataType.Image),
RandCropByPosNegLabeld(
keys=[DataType.Image, DataType.Label],
label_key=DataType.Label,
spatial_size=(
110, 110, 25
), # Crop to slightly larger than desired for elastic deformation - minimises errors created from padding
pos=1,
neg=0,
num_samples=1,
image_key=DataType.Image,
image_threshold=0,
),
# This function handles both affine and elastic deformation together
# To minimise padding issues, we will elastic and affine transform first on original before any cropping
Rand3DElasticd(
keys=[DataType.Image, DataType.Label],
sigma_range=(5, 8), # Sigma range for elastic deformation
magnitude_range=(100, 200), # maginitude range for elastic deformation
mode=('bilinear', 'nearest'),
# Output to desired crop size
spatial_size=(96, 96, 16),
# Probability of augmentation - we will kepp most unchanged and augment 30%
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 run_training_test(root_dir, device="cuda:0", cachedataset=0):
monai.config.print_config()
images = sorted(glob(os.path.join(root_dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
train_files = [{"img": img, "seg": seg} for img, seg in zip(images[:20], segs[:20])]
val_files = [{"img": img, "seg": seg} for img, seg in zip(images[-20:], segs[-20:])]
# define transforms for image and segmentation
train_transforms = Compose(
[
LoadImaged(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
ScaleIntensityd(keys="img"),
RandCropByPosNegLabeld(
keys=["img", "seg"], label_key="seg", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
),
RandRotate90d(keys=["img", "seg"], prob=0.8, spatial_axes=[0, 2]),
ToTensord(keys=["img", "seg"]),
]
)
train_transforms.set_random_state(1234)
val_transforms = Compose(
[
LoadImaged(keys=["img", "seg"]),
AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
# resampling with align_corners=True or dtype=float64 will generate
# slight different results between PyTorch 1.5 an 1.6
Spacingd(keys=["img", "seg"], pixdim=[1.2, 0.8, 0.7], mode=["bilinear", "nearest"], dtype=np.float32),
ScaleIntensityd(keys="img"),
ToTensord(keys=["img", "seg"]),
]
)
# create a training data loader
if cachedataset == 2:
train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.8)
elif cachedataset == 3:
train_ds = monai.data.LMDBDataset(data=train_files, transform=train_transforms)
else:
train_ds = monai.data.Dataset(data=train_files, transform=train_transforms)
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = monai.data.DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)
# create a validation data loader
val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
dice_metric = DiceMetric(include_background=True, reduction="mean")
# create UNet, DiceLoss and Adam optimizer
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(), 5e-4)
# start a typical PyTorch training
val_interval = 2
best_metric, best_metric_epoch = -1, -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter(log_dir=os.path.join(root_dir, "runs"))
model_filename = os.path.join(root_dir, "best_metric_model.pth")
for epoch in range(6):
print("-" * 10)
print(f"Epoch {epoch + 1}/{6}")
model.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs, labels = batch_data["img"].to(device), batch_data["seg"].to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_len = len(train_ds) // train_loader.batch_size
print(f"{step}/{epoch_len}, train_loss:{loss.item():0.4f}")
writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch +1} average loss:{epoch_loss:0.4f}")
if (epoch + 1) % val_interval == 0:
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["img"].to(device), val_data["seg"].to(device)
sw_batch_size, roi_size = 4, (96, 96, 96)
val_outputs = val_post_tran(sliding_window_inference(val_images, roi_size, sw_batch_size, model))
value, not_nans = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += not_nans.item()
metric_sum += value.item() * not_nans.item()
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(), model_filename)
print("saved new best metric model")
print(
f"current epoch {epoch +1} current mean dice: {metric:0.4f} "
f"best mean dice: {best_metric:0.4f} at epoch {best_metric_epoch}"
)
writer.add_scalar("val_mean_dice", metric, epoch + 1)
# plot the last model output as GIF image in TensorBoard with the corresponding image and label
plot_2d_or_3d_image(val_images, epoch + 1, writer, index=0, tag="image")
plot_2d_or_3d_image(val_labels, epoch + 1, writer, index=0, tag="label")
plot_2d_or_3d_image(val_outputs, epoch + 1, writer, index=0, tag="output")
print(f"train completed, best_metric: {best_metric:0.4f} at epoch: {best_metric_epoch}")
writer.close()
return epoch_loss_values, best_metric, best_metric_epoch
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 except 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([
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]),
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=2,
pin_memory=True,
sampler=train_sampler,
multiprocessing_context=multiprocessing_context,
)
# 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)
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 get_task_transforms(mode, task_id, pos_sample_num, neg_sample_num,
num_samples):
if mode != "test":
keys = ["image", "label"]
else:
keys = ["image"]
load_transforms = [
LoadImaged(keys=keys),
EnsureChannelFirstd(keys=keys),
]
# 2. sampling
sample_transforms = [
PreprocessAnisotropic(
keys=keys,
clip_values=clip_values[task_id],
pixdim=spacing[task_id],
normalize_values=normalize_values[task_id],
model_mode=mode,
),
]
# 3. spatial transforms
if mode == "train":
other_transforms = [
SpatialPadd(keys=["image", "label"],
spatial_size=patch_size[task_id]),
RandCropByPosNegLabeld(
keys=["image", "label"],
label_key="label",
spatial_size=patch_size[task_id],
pos=pos_sample_num,
neg=neg_sample_num,
num_samples=num_samples,
image_key="image",
image_threshold=0,
),
RandZoomd(
keys=["image", "label"],
min_zoom=0.9,
max_zoom=1.2,
mode=("trilinear", "nearest"),
align_corners=(True, None),
prob=0.15,
),
RandGaussianNoised(keys=["image"], std=0.01, prob=0.15),
RandGaussianSmoothd(
keys=["image"],
sigma_x=(0.5, 1.15),
sigma_y=(0.5, 1.15),
sigma_z=(0.5, 1.15),
prob=0.15,
),
RandScaleIntensityd(keys=["image"], factors=0.3, prob=0.15),
RandFlipd(["image", "label"], spatial_axis=[0], prob=0.5),
RandFlipd(["image", "label"], spatial_axis=[1], prob=0.5),
RandFlipd(["image", "label"], spatial_axis=[2], prob=0.5),
CastToTyped(keys=["image", "label"], dtype=(np.float32, np.uint8)),
EnsureTyped(keys=["image", "label"]),
]
elif mode == "validation":
other_transforms = [
CastToTyped(keys=["image", "label"], dtype=(np.float32, np.uint8)),
EnsureTyped(keys=["image", "label"]),
]
else:
other_transforms = [
CastToTyped(keys=["image"], dtype=(np.float32)),
EnsureTyped(keys=["image"]),
]
all_transforms = load_transforms + sample_transforms + other_transforms
return Compose(all_transforms)
def main():
print_config()
# Define paths for running the script
data_dir = os.path.normpath('/to/be/defined')
json_path = os.path.normpath('/to/be/defined')
logdir = os.path.normpath('/to/be/defined')
# If use_pretrained is set to 0, ViT weights will not be loaded and random initialization is used
use_pretrained = 1
pretrained_path = os.path.normpath('/to/be/defined')
# Training Hyper-parameters
lr = 1e-4
max_iterations = 30000
eval_num = 100
if os.path.exists(logdir) == False:
os.mkdir(logdir)
# Training & Validation Transform chain
train_transforms = Compose([
LoadImaged(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Spacingd(
keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest"),
),
Orientationd(keys=["image", "label"], axcodes="RAS"),
ScaleIntensityRanged(
keys=["image"],
a_min=-175,
a_max=250,
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,
),
RandFlipd(
keys=["image", "label"],
spatial_axis=[0],
prob=0.10,
),
RandFlipd(
keys=["image", "label"],
spatial_axis=[1],
prob=0.10,
),
RandFlipd(
keys=["image", "label"],
spatial_axis=[2],
prob=0.10,
),
RandRotate90d(
keys=["image", "label"],
prob=0.10,
max_k=3,
),
RandShiftIntensityd(
keys=["image"],
offsets=0.10,
prob=0.50,
),
ToTensord(keys=["image", "label"]),
])
val_transforms = Compose([
LoadImaged(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Spacingd(
keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest"),
),
Orientationd(keys=["image", "label"], axcodes="RAS"),
ScaleIntensityRanged(keys=["image"],
a_min=-175,
a_max=250,
b_min=0.0,
b_max=1.0,
clip=True),
CropForegroundd(keys=["image", "label"], source_key="image"),
ToTensord(keys=["image", "label"]),
])
datalist = load_decathlon_datalist(base_dir=data_dir,
data_list_file_path=json_path,
is_segmentation=True,
data_list_key="training")
val_files = load_decathlon_datalist(base_dir=data_dir,
data_list_file_path=json_path,
is_segmentation=True,
data_list_key="validation")
train_ds = CacheDataset(
data=datalist,
transform=train_transforms,
cache_num=24,
cache_rate=1.0,
num_workers=4,
)
train_loader = DataLoader(train_ds,
batch_size=1,
shuffle=True,
num_workers=4,
pin_memory=True)
val_ds = CacheDataset(data=val_files,
transform=val_transforms,
cache_num=6,
cache_rate=1.0,
num_workers=4)
val_loader = DataLoader(val_ds,
batch_size=1,
shuffle=False,
num_workers=4,
pin_memory=True)
case_num = 0
img = val_ds[case_num]["image"]
label = val_ds[case_num]["label"]
img_shape = img.shape
label_shape = label.shape
print(f"image shape: {img_shape}, label shape: {label_shape}")
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNETR(
in_channels=1,
out_channels=14,
img_size=(96, 96, 96),
feature_size=16,
hidden_size=768,
mlp_dim=3072,
num_heads=12,
pos_embed="conv",
norm_name="instance",
res_block=True,
dropout_rate=0.0,
)
# Load ViT backbone weights into UNETR
if use_pretrained == 1:
print('Loading Weights from the Path {}'.format(pretrained_path))
vit_dict = torch.load(pretrained_path)
vit_weights = vit_dict['state_dict']
# Delete the following variable names conv3d_transpose.weight, conv3d_transpose.bias,
# conv3d_transpose_1.weight, conv3d_transpose_1.bias as they were used to match dimensions
# while pretraining with ViTAutoEnc and are not a part of ViT backbone (this is used in UNETR)
vit_weights.pop('conv3d_transpose_1.bias')
vit_weights.pop('conv3d_transpose_1.weight')
vit_weights.pop('conv3d_transpose.bias')
vit_weights.pop('conv3d_transpose.weight')
model.vit.load_state_dict(vit_weights)
print('Pretrained Weights Succesfully Loaded !')
elif use_pretrained == 0:
print(
'No weights were loaded, all weights being used are randomly initialized!'
)
model.to(device)
loss_function = DiceCELoss(to_onehot_y=True, softmax=True)
torch.backends.cudnn.benchmark = True
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)
post_label = AsDiscrete(to_onehot=14)
post_pred = AsDiscrete(argmax=True, to_onehot=14)
dice_metric = DiceMetric(include_background=True,
reduction="mean",
get_not_nans=False)
global_step = 0
dice_val_best = 0.0
global_step_best = 0
epoch_loss_values = []
metric_values = []
def validation(epoch_iterator_val):
model.eval()
dice_vals = list()
with torch.no_grad():
for step, batch in enumerate(epoch_iterator_val):
val_inputs, val_labels = (batch["image"].cuda(),
batch["label"].cuda())
val_outputs = sliding_window_inference(val_inputs,
(96, 96, 96), 4, model)
val_labels_list = decollate_batch(val_labels)
val_labels_convert = [
post_label(val_label_tensor)
for val_label_tensor in val_labels_list
]
val_outputs_list = decollate_batch(val_outputs)
val_output_convert = [
post_pred(val_pred_tensor)
for val_pred_tensor in val_outputs_list
]
dice_metric(y_pred=val_output_convert, y=val_labels_convert)
dice = dice_metric.aggregate().item()
dice_vals.append(dice)
epoch_iterator_val.set_description(
"Validate (%d / %d Steps) (dice=%2.5f)" %
(global_step, 10.0, dice))
dice_metric.reset()
mean_dice_val = np.mean(dice_vals)
return mean_dice_val
def train(global_step, train_loader, dice_val_best, global_step_best):
model.train()
epoch_loss = 0
step = 0
epoch_iterator = tqdm(train_loader,
desc="Training (X / X Steps) (loss=X.X)",
dynamic_ncols=True)
for step, batch in enumerate(epoch_iterator):
step += 1
x, y = (batch["image"].cuda(), batch["label"].cuda())
logit_map = model(x)
loss = loss_function(logit_map, y)
loss.backward()
epoch_loss += loss.item()
optimizer.step()
optimizer.zero_grad()
epoch_iterator.set_description(
"Training (%d / %d Steps) (loss=%2.5f)" %
(global_step, max_iterations, loss))
if (global_step % eval_num == 0
and global_step != 0) or global_step == max_iterations:
epoch_iterator_val = tqdm(
val_loader,
desc="Validate (X / X Steps) (dice=X.X)",
dynamic_ncols=True)
dice_val = validation(epoch_iterator_val)
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
metric_values.append(dice_val)
if dice_val > dice_val_best:
dice_val_best = dice_val
global_step_best = global_step
torch.save(model.state_dict(),
os.path.join(logdir, "best_metric_model.pth"))
print(
"Model Was Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}"
.format(dice_val_best, dice_val))
else:
print(
"Model Was Not Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}"
.format(dice_val_best, dice_val))
plt.figure(1, (12, 6))
plt.subplot(1, 2, 1)
plt.title("Iteration Average Loss")
x = [eval_num * (i + 1) for i in range(len(epoch_loss_values))]
y = epoch_loss_values
plt.xlabel("Iteration")
plt.plot(x, y)
plt.grid()
plt.subplot(1, 2, 2)
plt.title("Val Mean Dice")
x = [eval_num * (i + 1) for i in range(len(metric_values))]
y = metric_values
plt.xlabel("Iteration")
plt.plot(x, y)
plt.grid()
plt.savefig(
os.path.join(logdir, 'btcv_finetune_quick_update.png'))
plt.clf()
plt.close(1)
global_step += 1
return global_step, dice_val_best, global_step_best
while global_step < max_iterations:
global_step, dice_val_best, global_step_best = train(
global_step, train_loader, dice_val_best, global_step_best)
model.load_state_dict(
torch.load(os.path.join(logdir, "best_metric_model.pth")))
print(f"train completed, best_metric: {dice_val_best:.4f} "
f"at iteration: {global_step_best}")
def train(args):
# disable logging for processes execpt 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([
LoadNiftid(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]),
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)
# 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}")
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)
# wrap the model with DistributedDataParallel module
model = DistributedDataParallel(model, device_ids=[args.local_rank])
# 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()
'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',
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
RandAffined(KEYS, spatial_size=None, prob=0)))
TESTS.append((
"RandCropByLabelClassesd 2d",
"2D",
1e-7,
True,
RandCropByLabelClassesd(KEYS,
"label", (99, 96),
ratios=[1, 2, 3, 4, 5],
num_classes=5,
num_samples=10),
))
TESTS.append(("RandCropByPosNegLabeld 2d", "2D", 1e-7, True,
RandCropByPosNegLabeld(KEYS, "label", (99, 96), num_samples=10)))
TESTS.append(("RandSpatialCropSamplesd 2d", "2D", 1e-7, True,
RandSpatialCropSamplesd(KEYS, (90, 91), num_samples=10)))
TESTS.append(("RandWeightedCropd 2d", "2D", 1e-7, True,
RandWeightedCropd(KEYS, "label", (90, 91), num_samples=10)))
TESTS_COMPOSE_X2 = [(t[0] + " Compose", t[1], t[2], t[3],
Compose(Compose(t[4:]))) for t in TESTS]
TESTS = TESTS + TESTS_COMPOSE_X2 # type: ignore
NUM_SAMPLES = 5
N_SAMPLES_TESTS = [
[
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)
