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_image_scale_intensity_ranged_none(self):
key = "img"
scaler = ScaleIntensityRanged(keys=key, a_min=20, a_max=108, b_min=None, b_max=None)
for p in TEST_NDARRAYS:
scaled = scaler({key: p(self.imt)})
expected = (self.imt - 20) / 88
assert_allclose(scaled[key], p(expected))
def pre_transforms(self, data):
return [
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
AddBackgroundScribblesFromROId(
scribbles="label",
scribbles_bg_label=self.scribbles_bg_label,
scribbles_fg_label=self.scribbles_fg_label,
),
# at the moment optimisers are bottleneck taking a long time,
# therefore scaling non-isotropic with big spacing
Spacingd(keys=["image", "label"], pixdim=self.pix_dim, mode=["bilinear", "nearest"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
ScaleIntensityRanged(
keys="image",
a_min=self.intensity_range[0],
a_max=self.intensity_range[1],
b_min=self.intensity_range[2],
b_max=self.intensity_range[3],
clip=self.intensity_range[4],
),
MakeLikelihoodFromScribblesHistogramd(
image="image",
scribbles="label",
post_proc_label="prob",
scribbles_bg_label=self.scribbles_bg_label,
scribbles_fg_label=self.scribbles_fg_label,
normalise=True,
),
]
def val_pre_transforms(self, context: Context):
return [
LoadImaged(keys=("image", "label"), reader="ITKReader"),
NormalizeLabelsInDatasetd(keys="label", label_names=self._labels),
EnsureChannelFirstd(keys=("image", "label")),
Orientationd(keys=["image", "label"], axcodes="RAS"),
# This transform may not work well for MR images
ScaleIntensityRanged(keys=("image"),
a_min=-175,
a_max=250,
b_min=0.0,
b_max=1.0,
clip=True),
Resized(keys=("image", "label"),
spatial_size=self.spatial_size,
mode=("area", "nearest")),
# Transforms for click simulation
FindAllValidSlicesMissingLabelsd(keys="label", sids="sids"),
AddInitialSeedPointMissingLabelsd(keys="label",
guidance="guidance",
sids="sids"),
AddGuidanceSignalCustomd(
keys="image",
guidance="guidance",
number_intensity_ch=self.number_intensity_ch),
#
ToTensord(keys=("image", "label")),
SelectItemsd(keys=("image", "label", "guidance", "label_names")),
]
def pre_transforms(self, data=None):
t = [
LoadImaged(keys="image", reader="ITKReader"),
EnsureChannelFirstd(keys="image"),
Orientationd(keys="image", axcodes="RAS"),
ScaleIntensityRanged(keys="image", a_min=-175, a_max=250, b_min=0.0, b_max=1.0, clip=True),
]
if self.type == InferType.DEEPEDIT:
t.extend(
[
AddGuidanceFromPointsCustomd(ref_image="image", guidance="guidance", label_names=self.labels),
Resized(keys="image", spatial_size=self.spatial_size, mode="area"),
ResizeGuidanceMultipleLabelCustomd(guidance="guidance", ref_image="image"),
AddGuidanceSignalCustomd(
keys="image", guidance="guidance", number_intensity_ch=self.number_intensity_ch
),
]
)
else:
t.extend(
[
Resized(keys="image", spatial_size=self.spatial_size, mode="area"),
DiscardAddGuidanced(
keys="image", label_names=self.labels, number_intensity_ch=self.number_intensity_ch
),
]
)
t.append(EnsureTyped(keys="image", device=data.get("device") if data else None))
return t
def get_xforms_load(mode="load", keys=("image", "label")):
"""returns a composed transform."""
xforms = [
LoadImaged(keys),
ScaleIntensityRanged(keys[0], a_min=-1000.0, a_max=500.0, b_min=0.0, b_max=1.0, clip=True),
]
if mode == "load":
dtype = (np.float32, np.uint8)
xforms.extend([CastToTyped(keys, dtype=dtype)])
return monai.transforms.Compose(xforms)
def test_image_scale_intensity_ranged(self):
key = "img"
scaler = ScaleIntensityRanged(keys=key,
a_min=20,
a_max=108,
b_min=50,
b_max=80)
scaled = scaler({key: self.imt})
expected = (self.imt - 20) / 88
expected = expected * 30 + 50
self.assertTrue(np.allclose(scaled[key], expected))
def test_image_scale_intensity_ranged(self):
key = "img"
scaler = ScaleIntensityRanged(keys=key,
a_min=20,
a_max=108,
b_min=50,
b_max=80)
for p in TEST_NDARRAYS:
scaled = scaler({key: p(self.imt)})
expected = (self.imt - 20) / 88
expected = expected * 30 + 50
assert_allclose(scaled[key], p(expected), type_test="tensor")
def pre_transforms(self, data=None) -> Sequence[Callable]:
return [
LoadImaged(keys="image"),
AddChanneld(keys="image"),
Spacingd(keys="image", pixdim=[1.0, 1.0, 1.0]),
ScaleIntensityRanged(keys="image",
a_min=-57,
a_max=164,
b_min=0.0,
b_max=1.0,
clip=True),
EnsureTyped(keys="image"),
]
def pre_transforms(self, data=None) -> Sequence[Callable]:
return [
LoadImaged(keys="image", reader="ITKReader"),
EnsureChannelFirstd(keys="image"),
Spacingd(keys="image", pixdim=self.target_spacing),
ScaleIntensityRanged(keys="image",
a_min=-175,
a_max=250,
b_min=0.0,
b_max=1.0,
clip=True),
EnsureTyped(keys="image"),
]
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 val_pre_transforms(self, context: Context):
return [
LoadImaged(keys=("image", "label"), reader="ITKReader"),
EnsureChannelFirstd(keys=("image", "label")),
Spacingd(keys=("image", "label"),
pixdim=self.target_spacing,
mode=("bilinear", "nearest")),
ScaleIntensityRanged(keys="image",
a_min=-175,
a_max=250,
b_min=0.0,
b_max=1.0,
clip=True),
EnsureTyped(keys=("image", "label")),
SelectItemsd(keys=("image", "label")),
]
def val_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),
]
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_xforms_scans_or_synthetic_lesions(mode="scans",
keys=("image", "label")):
"""returns a composed transform for scans or synthetic lesions."""
xforms = [
LoadImaged(keys),
AddChanneld(keys),
Orientationd(keys, axcodes="LPS"),
Spacingd(keys,
pixdim=(1.25, 1.25, 5.0),
mode=("bilinear", "nearest")[:len(keys)]),
]
dtype = (np.int16, np.uint8)
if mode == "synthetic":
xforms.extend([
ScaleIntensityRanged(keys[0],
a_min=-1000.0,
a_max=500.0,
b_min=0.0,
b_max=1.0,
clip=True),
])
dtype = (np.float32, np.uint8)
xforms.extend([CastToTyped(keys, dtype=dtype)])
return monai.transforms.Compose(xforms)
def default_transform(self, train=False):
target_shape = (128, 128, 256) # (x, y, z)
target_voxel_spacing = (4.8, 4.8, 4.8) # (x, y, z)
if train:
return Compose([ # read img + meta info
LoadNifti(keys=["pet_img", "ct_img", "mask_img"]),
Roi2Mask(keys=['pet_img', 'mask_img'],
method='otsu',
tval=0.0,
idx_channel=0),
ResampleReshapeAlign(target_shape,
target_voxel_spacing,
keys=['pet_img', "ct_img", 'mask_img'],
origin='head',
origin_key='pet_img'),
Sitk2Numpy(keys=['pet_img', 'ct_img', 'mask_img']),
# user can also add other random transforms
RandAffined(keys=("pet_img", "ct_img", "mask_img"),
spatial_size=None,
prob=0.4,
rotate_range=(0, np.pi / 30, np.pi / 15),
shear_range=None,
translate_range=(10, 10, 10),
scale_range=(0.1, 0.1, 0.1),
mode=("bilinear", "bilinear", "nearest"),
padding_mode="border"),
# semantic to instance segmentation
ConnectedComponent(keys='mask_img',
channels_first=True,
exclude_background=True),
# Add bouding Box
GenerateBbox(keys='mask_img'),
# normalize input
ScaleIntensityRanged(
keys=["pet_img"],
a_min=0.0,
a_max=25.0,
b_min=0.0,
b_max=1.0,
clip=True,
),
ScaleIntensityRanged(
keys=["ct_img"],
a_min=-1000.0,
a_max=1000.0,
b_min=0.0,
b_max=1.0,
clip=True,
),
# Prepare for neural network
ConcatModality(keys=['pet_img', 'ct_img']),
ToTensord(keys=["image", "mask_img"]),
])
else:
return Compose([ # read img + meta info
LoadNifti(keys=["pet_img", "ct_img", "mask_img"]),
Roi2Mask(keys=['pet_img', 'mask_img'],
method='otsu',
tval=0.0,
idx_channel=0),
ResampleReshapeAlign(target_shape,
target_voxel_spacing,
keys=['pet_img', "ct_img", 'mask_img'],
origin='head',
origin_key='pet_img'),
Sitk2Numpy(keys=['pet_img', 'ct_img', 'mask_img']),
# semantic to instance segmentation
ConnectedComponent(keys='mask_img',
channels_first=True,
exclude_background=True),
# Add bouding Box
GenerateBbox(keys='mask_img'),
# normalize input
ScaleIntensityRanged(
keys=["pet_img"],
a_min=0.0,
a_max=25.0,
b_min=0.0,
b_max=1.0,
clip=True,
),
ScaleIntensityRanged(
keys=["ct_img"],
a_min=-1000.0,
a_max=1000.0,
b_min=0.0,
b_max=1.0,
clip=True,
),
# Prepare for neural network
ConcatModality(keys=['pet_img', 'ct_img']),
ToTensord(keys=["image", "mask_img"]),
])
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)
[
Range("LoadImage")(LoadImaged(keys=["image", "label"])),
Range()(EnsureChannelFirstd(keys=["image", "label"])),
Range("Spacing")(
Spacingd(
keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest"),
)
),
Range()(Orientationd(keys=["image", "label"], axcodes="RAS")),
Range()(
ScaleIntensityRanged(
keys=["image"],
a_min=-57,
a_max=164,
b_min=0.0,
b_max=1.0,
clip=True,
)
),
Range()(CropForegroundd(keys=["image", "label"], source_key="image")),
# pre-compute foreground and background indexes
# and cache them to accelerate training
Range("Indexing")(
FgBgToIndicesd(
keys="label",
fg_postfix="_fg",
bg_postfix="_bg",
image_key="image",
)
),
'imagesTr'),
os.path.join(params['data_root'],
'labelsTr'),
file_pattern='*.nii.gz')
set_determinism(seed=0)
val_transforms = Compose([
LoadNiftid(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
Spacingd(keys=['image', 'label'],
pixdim=(1.5, 1.5, 2.),
mode=('bilinear', 'nearest')),
Orientationd(keys=['image', 'label'], axcodes='RAS'),
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'])
])
check_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
check_loader = monai.data.DataLoader(check_ds, batch_size=1)
check_data = monai.utils.misc.first(check_loader)
# plt.imshow(check_data['image'][0, 0, :, :, 80])
# plt.imshow(check_data['label'][0, 0, :, :, 80])
val_ds = monai.data.CacheDataset(data=val_files,
transform=val_transforms,
cache_rate=1.0,
def main():
parser = argparse.ArgumentParser(description="training")
parser.add_argument(
"--checkpoint",
type=str,
default=None,
help="checkpoint full path",
)
parser.add_argument(
"--factor_ram_cost",
default=0.0,
type=float,
help="factor to determine RAM cost in the searched architecture",
)
parser.add_argument(
"--fold",
action="store",
required=True,
help="fold index in N-fold cross-validation",
)
parser.add_argument(
"--json",
action="store",
required=True,
help="full path of .json file",
)
parser.add_argument(
"--json_key",
action="store",
required=True,
help="selected key in .json data list",
)
parser.add_argument(
"--local_rank",
required=int,
help="local process rank",
)
parser.add_argument(
"--num_folds",
action="store",
required=True,
help="number of folds in cross-validation",
)
parser.add_argument(
"--output_root",
action="store",
required=True,
help="output root",
)
parser.add_argument(
"--root",
action="store",
required=True,
help="data root",
)
args = parser.parse_args()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
if not os.path.exists(args.output_root):
os.makedirs(args.output_root, exist_ok=True)
amp = True
determ = True
factor_ram_cost = args.factor_ram_cost
fold = int(args.fold)
input_channels = 1
learning_rate = 0.025
learning_rate_arch = 0.001
learning_rate_milestones = np.array([0.4, 0.8])
num_images_per_batch = 1
num_epochs = 1430 # around 20k iteration
num_epochs_per_validation = 100
num_epochs_warmup = 715
num_folds = int(args.num_folds)
num_patches_per_image = 1
num_sw_batch_size = 6
output_classes = 3
overlap_ratio = 0.625
patch_size = (96, 96, 96)
patch_size_valid = (96, 96, 96)
spacing = [1.0, 1.0, 1.0]
print("factor_ram_cost", factor_ram_cost)
# deterministic training
if determ:
set_determinism(seed=0)
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
# dist.barrier()
world_size = dist.get_world_size()
with open(args.json, "r") as f:
json_data = json.load(f)
split = len(json_data[args.json_key]) // num_folds
list_train = json_data[args.json_key][:(
split * fold)] + json_data[args.json_key][(split * (fold + 1)):]
list_valid = json_data[args.json_key][(split * fold):(split * (fold + 1))]
# training data
files = []
for _i in range(len(list_train)):
str_img = os.path.join(args.root, list_train[_i]["image"])
str_seg = os.path.join(args.root, list_train[_i]["label"])
if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
continue
files.append({"image": str_img, "label": str_seg})
train_files = files
random.shuffle(train_files)
train_files_w = train_files[:len(train_files) // 2]
train_files_w = partition_dataset(data=train_files_w,
shuffle=True,
num_partitions=world_size,
even_divisible=True)[dist.get_rank()]
print("train_files_w:", len(train_files_w))
train_files_a = train_files[len(train_files) // 2:]
train_files_a = partition_dataset(data=train_files_a,
shuffle=True,
num_partitions=world_size,
even_divisible=True)[dist.get_rank()]
print("train_files_a:", len(train_files_a))
# validation data
files = []
for _i in range(len(list_valid)):
str_img = os.path.join(args.root, list_valid[_i]["image"])
str_seg = os.path.join(args.root, list_valid[_i]["label"])
if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
continue
files.append({"image": str_img, "label": str_seg})
val_files = files
val_files = partition_dataset(data=val_files,
shuffle=False,
num_partitions=world_size,
even_divisible=False)[dist.get_rank()]
print("val_files:", len(val_files))
# network architecture
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
train_transforms = Compose([
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(keys=["image", "label"],
pixdim=spacing,
mode=("bilinear", "nearest"),
align_corners=(True, True)),
CastToTyped(keys=["image"], dtype=(torch.float32)),
ScaleIntensityRanged(keys=["image"],
a_min=-87.0,
a_max=199.0,
b_min=0.0,
b_max=1.0,
clip=True),
CastToTyped(keys=["image", "label"], dtype=(np.float16, np.uint8)),
CopyItemsd(keys=["label"], times=1, names=["label4crop"]),
Lambdad(
keys=["label4crop"],
func=lambda x: np.concatenate(tuple([
ndimage.binary_dilation(
(x == _k).astype(x.dtype), iterations=48).astype(x.dtype)
for _k in range(output_classes)
]),
axis=0),
overwrite=True,
),
EnsureTyped(keys=["image", "label"]),
CastToTyped(keys=["image"], dtype=(torch.float32)),
SpatialPadd(keys=["image", "label", "label4crop"],
spatial_size=patch_size,
mode=["reflect", "constant", "constant"]),
RandCropByLabelClassesd(keys=["image", "label"],
label_key="label4crop",
num_classes=output_classes,
ratios=[
1,
] * output_classes,
spatial_size=patch_size,
num_samples=num_patches_per_image),
Lambdad(keys=["label4crop"], func=lambda x: 0),
RandRotated(keys=["image", "label"],
range_x=0.3,
range_y=0.3,
range_z=0.3,
mode=["bilinear", "nearest"],
prob=0.2),
RandZoomd(keys=["image", "label"],
min_zoom=0.8,
max_zoom=1.2,
mode=["trilinear", "nearest"],
align_corners=[True, None],
prob=0.16),
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.5),
RandShiftIntensityd(keys=["image"], offsets=0.1, prob=0.5),
RandGaussianNoised(keys=["image"], std=0.01, prob=0.15),
RandFlipd(keys=["image", "label"], spatial_axis=0, prob=0.5),
RandFlipd(keys=["image", "label"], spatial_axis=1, prob=0.5),
RandFlipd(keys=["image", "label"], spatial_axis=2, prob=0.5),
CastToTyped(keys=["image", "label"],
dtype=(torch.float32, torch.uint8)),
ToTensord(keys=["image", "label"]),
])
val_transforms = Compose([
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(keys=["image", "label"],
pixdim=spacing,
mode=("bilinear", "nearest"),
align_corners=(True, True)),
CastToTyped(keys=["image"], dtype=(torch.float32)),
ScaleIntensityRanged(keys=["image"],
a_min=-87.0,
a_max=199.0,
b_min=0.0,
b_max=1.0,
clip=True),
CastToTyped(keys=["image", "label"], dtype=(np.float32, np.uint8)),
EnsureTyped(keys=["image", "label"]),
ToTensord(keys=["image", "label"])
])
train_ds_a = monai.data.CacheDataset(data=train_files_a,
transform=train_transforms,
cache_rate=1.0,
num_workers=8)
train_ds_w = monai.data.CacheDataset(data=train_files_w,
transform=train_transforms,
cache_rate=1.0,
num_workers=8)
val_ds = monai.data.CacheDataset(data=val_files,
transform=val_transforms,
cache_rate=1.0,
num_workers=2)
# monai.data.Dataset can be used as alternatives when debugging or RAM space is limited.
# train_ds_a = monai.data.Dataset(data=train_files_a, transform=train_transforms)
# train_ds_w = monai.data.Dataset(data=train_files_w, transform=train_transforms)
# val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
train_loader_a = ThreadDataLoader(train_ds_a,
num_workers=0,
batch_size=num_images_per_batch,
shuffle=True)
train_loader_w = ThreadDataLoader(train_ds_w,
num_workers=0,
batch_size=num_images_per_batch,
shuffle=True)
val_loader = ThreadDataLoader(val_ds,
num_workers=0,
batch_size=1,
shuffle=False)
# DataLoader can be used as alternatives when ThreadDataLoader is less efficient.
# train_loader_a = DataLoader(train_ds_a, batch_size=num_images_per_batch, shuffle=True, num_workers=2, pin_memory=torch.cuda.is_available())
# train_loader_w = DataLoader(train_ds_w, batch_size=num_images_per_batch, shuffle=True, num_workers=2, pin_memory=torch.cuda.is_available())
# val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=2, pin_memory=torch.cuda.is_available())
dints_space = monai.networks.nets.TopologySearch(
channel_mul=0.5,
num_blocks=12,
num_depths=4,
use_downsample=True,
device=device,
)
model = monai.networks.nets.DiNTS(
dints_space=dints_space,
in_channels=input_channels,
num_classes=output_classes,
use_downsample=True,
)
model = model.to(device)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
post_pred = Compose(
[EnsureType(),
AsDiscrete(argmax=True, to_onehot=output_classes)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=output_classes)])
# loss function
loss_func = monai.losses.DiceCELoss(
include_background=False,
to_onehot_y=True,
softmax=True,
squared_pred=True,
batch=True,
smooth_nr=0.00001,
smooth_dr=0.00001,
)
# optimizer
optimizer = torch.optim.SGD(model.weight_parameters(),
lr=learning_rate * world_size,
momentum=0.9,
weight_decay=0.00004)
arch_optimizer_a = torch.optim.Adam([dints_space.log_alpha_a],
lr=learning_rate_arch * world_size,
betas=(0.5, 0.999),
weight_decay=0.0)
arch_optimizer_c = torch.optim.Adam([dints_space.log_alpha_c],
lr=learning_rate_arch * world_size,
betas=(0.5, 0.999),
weight_decay=0.0)
print()
if torch.cuda.device_count() > 1:
if dist.get_rank() == 0:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = DistributedDataParallel(model,
device_ids=[device],
find_unused_parameters=True)
if args.checkpoint != None and os.path.isfile(args.checkpoint):
print("[info] fine-tuning pre-trained checkpoint {0:s}".format(
args.checkpoint))
model.load_state_dict(torch.load(args.checkpoint, map_location=device))
torch.cuda.empty_cache()
else:
print("[info] training from scratch")
# amp
if amp:
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
if dist.get_rank() == 0:
print("[info] amp enabled")
# start a typical PyTorch training
val_interval = num_epochs_per_validation
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
idx_iter = 0
metric_values = list()
if dist.get_rank() == 0:
writer = SummaryWriter(
log_dir=os.path.join(args.output_root, "Events"))
with open(os.path.join(args.output_root, "accuracy_history.csv"),
"a") as f:
f.write("epoch\tmetric\tloss\tlr\ttime\titer\n")
dataloader_a_iterator = iter(train_loader_a)
start_time = time.time()
for epoch in range(num_epochs):
decay = 0.5**np.sum([
(epoch - num_epochs_warmup) /
(num_epochs - num_epochs_warmup) > learning_rate_milestones
])
lr = learning_rate * decay
for param_group in optimizer.param_groups:
param_group["lr"] = lr
if dist.get_rank() == 0:
print("-" * 10)
print(f"epoch {epoch + 1}/{num_epochs}")
print("learning rate is set to {}".format(lr))
model.train()
epoch_loss = 0
loss_torch = torch.zeros(2, dtype=torch.float, device=device)
epoch_loss_arch = 0
loss_torch_arch = torch.zeros(2, dtype=torch.float, device=device)
step = 0
for batch_data in train_loader_w:
step += 1
inputs, labels = batch_data["image"].to(
device), batch_data["label"].to(device)
if world_size == 1:
for _ in model.weight_parameters():
_.requires_grad = True
else:
for _ in model.module.weight_parameters():
_.requires_grad = True
dints_space.log_alpha_a.requires_grad = False
dints_space.log_alpha_c.requires_grad = False
optimizer.zero_grad()
if amp:
with autocast():
outputs = model(inputs)
if output_classes == 2:
loss = loss_func(torch.flip(outputs, dims=[1]),
1 - labels)
else:
loss = loss_func(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
outputs = model(inputs)
if output_classes == 2:
loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
else:
loss = loss_func(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
loss_torch[0] += loss.item()
loss_torch[1] += 1.0
epoch_len = len(train_loader_w)
idx_iter += 1
if dist.get_rank() == 0:
print("[{0}] ".format(str(datetime.now())[:19]) +
f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")
writer.add_scalar("train_loss", loss.item(),
epoch_len * epoch + step)
if epoch < num_epochs_warmup:
continue
try:
sample_a = next(dataloader_a_iterator)
except StopIteration:
dataloader_a_iterator = iter(train_loader_a)
sample_a = next(dataloader_a_iterator)
inputs_search, labels_search = sample_a["image"].to(
device), sample_a["label"].to(device)
if world_size == 1:
for _ in model.weight_parameters():
_.requires_grad = False
else:
for _ in model.module.weight_parameters():
_.requires_grad = False
dints_space.log_alpha_a.requires_grad = True
dints_space.log_alpha_c.requires_grad = True
# linear increase topology and RAM loss
entropy_alpha_c = torch.tensor(0.).to(device)
entropy_alpha_a = torch.tensor(0.).to(device)
ram_cost_full = torch.tensor(0.).to(device)
ram_cost_usage = torch.tensor(0.).to(device)
ram_cost_loss = torch.tensor(0.).to(device)
topology_loss = torch.tensor(0.).to(device)
probs_a, arch_code_prob_a = dints_space.get_prob_a(child=True)
entropy_alpha_a = -((probs_a) * torch.log(probs_a + 1e-5)).mean()
entropy_alpha_c = -(F.softmax(dints_space.log_alpha_c, dim=-1) * \
F.log_softmax(dints_space.log_alpha_c, dim=-1)).mean()
topology_loss = dints_space.get_topology_entropy(probs_a)
ram_cost_full = dints_space.get_ram_cost_usage(inputs.shape,
full=True)
ram_cost_usage = dints_space.get_ram_cost_usage(inputs.shape)
ram_cost_loss = torch.abs(factor_ram_cost -
ram_cost_usage / ram_cost_full)
arch_optimizer_a.zero_grad()
arch_optimizer_c.zero_grad()
combination_weights = (epoch - num_epochs_warmup) / (
num_epochs - num_epochs_warmup)
if amp:
with autocast():
outputs_search = model(inputs_search)
if output_classes == 2:
loss = loss_func(torch.flip(outputs_search, dims=[1]),
1 - labels_search)
else:
loss = loss_func(outputs_search, labels_search)
loss += combination_weights * ((entropy_alpha_a + entropy_alpha_c) + ram_cost_loss \
+ 0.001 * topology_loss)
scaler.scale(loss).backward()
scaler.step(arch_optimizer_a)
scaler.step(arch_optimizer_c)
scaler.update()
else:
outputs_search = model(inputs_search)
if output_classes == 2:
loss = loss_func(torch.flip(outputs_search, dims=[1]),
1 - labels_search)
else:
loss = loss_func(outputs_search, labels_search)
loss += 1.0 * (combination_weights * (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss \
+ 0.001 * topology_loss)
loss.backward()
arch_optimizer_a.step()
arch_optimizer_c.step()
epoch_loss_arch += loss.item()
loss_torch_arch[0] += loss.item()
loss_torch_arch[1] += 1.0
if dist.get_rank() == 0:
print(
"[{0}] ".format(str(datetime.now())[:19]) +
f"{step}/{epoch_len}, train_loss_arch: {loss.item():.4f}")
writer.add_scalar("train_loss_arch", loss.item(),
epoch_len * epoch + step)
# synchronizes all processes and reduce results
dist.barrier()
dist.all_reduce(loss_torch, op=torch.distributed.ReduceOp.SUM)
loss_torch = loss_torch.tolist()
loss_torch_arch = loss_torch_arch.tolist()
if dist.get_rank() == 0:
loss_torch_epoch = loss_torch[0] / loss_torch[1]
print(
f"epoch {epoch + 1} average loss: {loss_torch_epoch:.4f}, best mean dice: {best_metric:.4f} at epoch {best_metric_epoch}"
)
if epoch >= num_epochs_warmup:
loss_torch_arch_epoch = loss_torch_arch[0] / loss_torch_arch[1]
print(
f"epoch {epoch + 1} average arch loss: {loss_torch_arch_epoch:.4f}, best mean dice: {best_metric:.4f} at epoch {best_metric_epoch}"
)
if (epoch + 1) % val_interval == 0:
torch.cuda.empty_cache()
model.eval()
with torch.no_grad():
metric = torch.zeros((output_classes - 1) * 2,
dtype=torch.float,
device=device)
metric_sum = 0.0
metric_count = 0
metric_mat = []
val_images = None
val_labels = None
val_outputs = None
_index = 0
for val_data in val_loader:
val_images = val_data["image"].to(device)
val_labels = val_data["label"].to(device)
roi_size = patch_size_valid
sw_batch_size = num_sw_batch_size
if amp:
with torch.cuda.amp.autocast():
pred = sliding_window_inference(
val_images,
roi_size,
sw_batch_size,
lambda x: model(x),
mode="gaussian",
overlap=overlap_ratio,
)
else:
pred = sliding_window_inference(
val_images,
roi_size,
sw_batch_size,
lambda x: model(x),
mode="gaussian",
overlap=overlap_ratio,
)
val_outputs = pred
val_outputs = post_pred(val_outputs[0, ...])
val_outputs = val_outputs[None, ...]
val_labels = post_label(val_labels[0, ...])
val_labels = val_labels[None, ...]
value = compute_meandice(y_pred=val_outputs,
y=val_labels,
include_background=False)
print(_index + 1, "/", len(val_loader), value)
metric_count += len(value)
metric_sum += value.sum().item()
metric_vals = value.cpu().numpy()
if len(metric_mat) == 0:
metric_mat = metric_vals
else:
metric_mat = np.concatenate((metric_mat, metric_vals),
axis=0)
for _c in range(output_classes - 1):
val0 = torch.nan_to_num(value[0, _c], nan=0.0)
val1 = 1.0 - torch.isnan(value[0, 0]).float()
metric[2 * _c] += val0 * val1
metric[2 * _c + 1] += val1
_index += 1
# synchronizes all processes and reduce results
dist.barrier()
dist.all_reduce(metric, op=torch.distributed.ReduceOp.SUM)
metric = metric.tolist()
if dist.get_rank() == 0:
for _c in range(output_classes - 1):
print(
"evaluation metric - class {0:d}:".format(_c + 1),
metric[2 * _c] / metric[2 * _c + 1])
avg_metric = 0
for _c in range(output_classes - 1):
avg_metric += metric[2 * _c] / metric[2 * _c + 1]
avg_metric = avg_metric / float(output_classes - 1)
print("avg_metric", avg_metric)
if avg_metric > best_metric:
best_metric = avg_metric
best_metric_epoch = epoch + 1
best_metric_iterations = idx_iter
node_a_d, arch_code_a_d, arch_code_c_d, arch_code_a_max_d = dints_space.decode(
)
torch.save(
{
"node_a": node_a_d,
"arch_code_a": arch_code_a_d,
"arch_code_a_max": arch_code_a_max_d,
"arch_code_c": arch_code_c_d,
"iter_num": idx_iter,
"epochs": epoch + 1,
"best_dsc": best_metric,
"best_path": best_metric_iterations,
},
os.path.join(args.output_root,
"search_code_" + str(idx_iter) + ".pth"),
)
print("saved new best metric model")
dict_file = {}
dict_file["best_avg_dice_score"] = float(best_metric)
dict_file["best_avg_dice_score_epoch"] = int(
best_metric_epoch)
dict_file["best_avg_dice_score_iteration"] = int(idx_iter)
with open(os.path.join(args.output_root, "progress.yaml"),
"w") as out_file:
documents = yaml.dump(dict_file, stream=out_file)
print(
"current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}"
.format(epoch + 1, avg_metric, best_metric,
best_metric_epoch))
current_time = time.time()
elapsed_time = (current_time - start_time) / 60.0
with open(
os.path.join(args.output_root,
"accuracy_history.csv"), "a") as f:
f.write(
"{0:d}\t{1:.5f}\t{2:.5f}\t{3:.5f}\t{4:.1f}\t{5:d}\n"
.format(epoch + 1, avg_metric, loss_torch_epoch,
lr, elapsed_time, idx_iter))
dist.barrier()
torch.cuda.empty_cache()
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
if dist.get_rank() == 0:
writer.close()
dist.destroy_process_group()
return
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 main():
opt = Options().parse()
# monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
if opt.gpu_ids != '-1':
num_gpus = len(opt.gpu_ids.split(','))
else:
num_gpus = 0
print('number of GPU:', num_gpus)
# Data loader creation
# train images
train_images = sorted(
glob(os.path.join(opt.images_folder, 'train', 'image*.nii')))
train_segs = sorted(
glob(os.path.join(opt.labels_folder, 'train', 'label*.nii')))
train_images_for_dice = sorted(
glob(os.path.join(opt.images_folder, 'train', 'image*.nii')))
train_segs_for_dice = sorted(
glob(os.path.join(opt.labels_folder, 'train', 'label*.nii')))
# validation images
val_images = sorted(
glob(os.path.join(opt.images_folder, 'val', 'image*.nii')))
val_segs = sorted(
glob(os.path.join(opt.labels_folder, 'val', 'label*.nii')))
# test images
test_images = sorted(
glob(os.path.join(opt.images_folder, 'test', 'image*.nii')))
test_segs = sorted(
glob(os.path.join(opt.labels_folder, 'test', 'label*.nii')))
# augment the data list for training
for i in range(int(opt.increase_factor_data)):
train_images.extend(train_images)
train_segs.extend(train_segs)
print('Number of training patches per epoch:', len(train_images))
print('Number of training images per epoch:', len(train_images_for_dice))
print('Number of validation images per epoch:', len(val_images))
print('Number of test images per epoch:', len(test_images))
# Creation of data directories for data_loader
train_dicts = [{
'image': image_name,
'label': label_name
} for image_name, label_name in zip(train_images, train_segs)]
train_dice_dicts = [{
'image': image_name,
'label': label_name
}
for image_name, label_name in zip(
train_images_for_dice, train_segs_for_dice)]
val_dicts = [{
'image': image_name,
'label': label_name
} for image_name, label_name in zip(val_images, val_segs)]
test_dicts = [{
'image': image_name,
'label': label_name
} for image_name, label_name in zip(test_images, test_segs)]
# Transforms list
if opt.resolution is not None:
train_transforms = [
LoadNiftid(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
ScaleIntensityRanged(
keys=["image"],
a_min=-120,
a_max=170,
b_min=0.0,
b_max=1.0,
clip=True,
),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
CropForegroundd(keys=["image", "label"], source_key="image"),
Spacingd(keys=['image', 'label'],
pixdim=opt.resolution,
mode=('bilinear', 'nearest')),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=1),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=0),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=2),
RandAffined(keys=['image', 'label'],
mode=('bilinear', 'nearest'),
prob=0.1,
rotate_range=(np.pi / 36, np.pi / 36, np.pi * 2),
padding_mode="zeros"),
RandAffined(keys=['image', 'label'],
mode=('bilinear', 'nearest'),
prob=0.1,
rotate_range=(np.pi / 36, np.pi / 2, np.pi / 36),
padding_mode="zeros"),
RandAffined(keys=['image', 'label'],
mode=('bilinear', 'nearest'),
prob=0.1,
rotate_range=(np.pi / 2, np.pi / 36, np.pi / 36),
padding_mode="zeros"),
Rand3DElasticd(keys=['image', 'label'],
mode=('bilinear', 'nearest'),
prob=0.1,
sigma_range=(5, 8),
magnitude_range=(100, 200),
scale_range=(0.15, 0.15, 0.15),
padding_mode="zeros"),
RandAdjustContrastd(keys=['image'], gamma=(0.5, 2.5), prob=0.1),
RandGaussianNoised(keys=['image'],
prob=0.1,
mean=np.random.uniform(0, 0.5),
std=np.random.uniform(0, 1)),
RandShiftIntensityd(keys=['image'],
offsets=np.random.uniform(0, 0.3),
prob=0.1),
RandSpatialCropd(keys=['image', 'label'],
roi_size=opt.patch_size,
random_size=False),
ToTensord(keys=['image', 'label'])
]
val_transforms = [
LoadNiftid(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
ScaleIntensityRanged(
keys=["image"],
a_min=-120,
a_max=170,
b_min=0.0,
b_max=1.0,
clip=True,
),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
CropForegroundd(keys=["image", "label"], source_key="image"),
Spacingd(keys=['image', 'label'],
pixdim=opt.resolution,
mode=('bilinear', 'nearest')),
ToTensord(keys=['image', 'label'])
]
else:
train_transforms = [
LoadNiftid(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
ScaleIntensityRanged(
keys=["image"],
a_min=-120,
a_max=170,
b_min=0.0,
b_max=1.0,
clip=True,
),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
CropForegroundd(keys=["image", "label"], source_key="image"),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=1),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=0),
RandFlipd(keys=['image', 'label'], prob=0.1, spatial_axis=2),
RandAffined(keys=['image', 'label'],
mode=('bilinear', 'nearest'),
prob=0.1,
rotate_range=(np.pi / 36, np.pi / 36, np.pi * 2),
padding_mode="zeros"),
RandAffined(keys=['image', 'label'],
mode=('bilinear', 'nearest'),
prob=0.1,
rotate_range=(np.pi / 36, np.pi / 2, np.pi / 36),
padding_mode="zeros"),
RandAffined(keys=['image', 'label'],
mode=('bilinear', 'nearest'),
prob=0.1,
rotate_range=(np.pi / 2, np.pi / 36, np.pi / 36),
padding_mode="zeros"),
Rand3DElasticd(keys=['image', 'label'],
mode=('bilinear', 'nearest'),
prob=0.1,
sigma_range=(5, 8),
magnitude_range=(100, 200),
scale_range=(0.15, 0.15, 0.15),
padding_mode="zeros"),
RandAdjustContrastd(keys=['image'], gamma=(0.5, 2.5), prob=0.1),
RandGaussianNoised(keys=['image'],
prob=0.1,
mean=np.random.uniform(0, 0.5),
std=np.random.uniform(0, 1)),
RandShiftIntensityd(keys=['image'],
offsets=np.random.uniform(0, 0.3),
prob=0.1),
RandSpatialCropd(keys=['image', 'label'],
roi_size=opt.patch_size,
random_size=False),
ToTensord(keys=['image', 'label'])
]
val_transforms = [
LoadNiftid(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
ScaleIntensityRanged(
keys=["image"],
a_min=-120,
a_max=170,
b_min=0.0,
b_max=1.0,
clip=True,
),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
CropForegroundd(keys=["image", "label"], source_key="image"),
ToTensord(keys=['image', 'label'])
]
train_transforms = Compose(train_transforms)
val_transforms = Compose(val_transforms)
# create a training data loader
check_train = monai.data.Dataset(data=train_dicts,
transform=train_transforms)
train_loader = DataLoader(check_train,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.workers,
pin_memory=torch.cuda.is_available())
# create a training_dice data loader
check_val = monai.data.Dataset(data=train_dice_dicts,
transform=val_transforms)
train_dice_loader = DataLoader(check_val,
batch_size=1,
num_workers=opt.workers,
pin_memory=torch.cuda.is_available())
# create a validation data loader
check_val = monai.data.Dataset(data=val_dicts, transform=val_transforms)
val_loader = DataLoader(check_val,
batch_size=1,
num_workers=opt.workers,
pin_memory=torch.cuda.is_available())
# create a validation data loader
check_val = monai.data.Dataset(data=test_dicts, transform=val_transforms)
test_loader = DataLoader(check_val,
batch_size=1,
num_workers=opt.workers,
pin_memory=torch.cuda.is_available())
# try to use all the available GPUs
devices = get_devices_spec(None)
# build the network
net = build_net()
net.cuda()
if num_gpus > 1:
net = torch.nn.DataParallel(net)
if opt.preload is not None:
net.load_state_dict(torch.load(opt.preload))
dice_metric = DiceMetric(include_background=True,
to_onehot_y=False,
sigmoid=True,
reduction="mean")
# loss_function = monai.losses.DiceLoss(sigmoid=True)
loss_function = monai.losses.TverskyLoss(sigmoid=True, alpha=0.3, beta=0.7)
optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
net_scheduler = get_scheduler(optim, opt)
# start a typical PyTorch training
val_interval = 1
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter()
for epoch in range(opt.epochs):
print("-" * 10)
print(f"epoch {epoch + 1}/{opt.epochs}")
net.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs, labels = batch_data["image"].cuda(
), batch_data["label"].cuda()
optim.zero_grad()
outputs = net(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optim.step()
epoch_loss += loss.item()
epoch_len = len(check_train) // 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}")
update_learning_rate(net_scheduler, optim)
if (epoch + 1) % val_interval == 0:
net.eval()
with torch.no_grad():
def plot_dice(images_loader):
metric_sum = 0.0
metric_count = 0
val_images = None
val_labels = None
val_outputs = None
for data in images_loader:
val_images, val_labels = data["image"].cuda(
), data["label"].cuda()
roi_size = opt.patch_size
sw_batch_size = 4
val_outputs = sliding_window_inference(
val_images, roi_size, sw_batch_size, net)
value = dice_metric(y_pred=val_outputs, y=val_labels)
metric_count += len(value)
metric_sum += value.item() * len(value)
metric = metric_sum / metric_count
metric_values.append(metric)
return metric, val_images, val_labels, val_outputs
metric, val_images, val_labels, val_outputs = plot_dice(
val_loader)
# Save best model
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
torch.save(net.state_dict(), "best_metric_model.pth")
print("saved new best metric model")
metric_train, train_images, train_labels, train_outputs = plot_dice(
train_dice_loader)
metric_test, test_images, test_labels, test_outputs = plot_dice(
test_loader)
# Logger bar
print(
"current epoch: {} Training dice: {:.4f} Validation dice: {:.4f} Testing dice: {:.4f} Best Validation dice: {:.4f} at epoch {}"
.format(epoch + 1, metric_train, metric, metric_test,
best_metric, best_metric_epoch))
writer.add_scalar("Mean_epoch_loss", epoch_loss, epoch + 1)
writer.add_scalar("Testing_dice", metric_test, epoch + 1)
writer.add_scalar("Training_dice", metric_train, epoch + 1)
writer.add_scalar("Validation_dice", metric, epoch + 1)
# plot the last model output as GIF image in TensorBoard with the corresponding image and label
val_outputs = (val_outputs.sigmoid() >= 0.5).float()
plot_2d_or_3d_image(val_images,
epoch + 1,
writer,
index=0,
tag="validation image")
plot_2d_or_3d_image(val_labels,
epoch + 1,
writer,
index=0,
tag="validation label")
plot_2d_or_3d_image(val_outputs,
epoch + 1,
writer,
index=0,
tag="validation inference")
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
writer.close()
def main():
#TODO Defining file paths & output directory path
json_Path = os.path.normpath('/scratch/data_2021/tcia_covid19/dataset_split_debug.json')
data_Root = os.path.normpath('/scratch/data_2021/tcia_covid19')
logdir_path = os.path.normpath('/home/vishwesh/monai_tutorial_testing/issue_467')
if os.path.exists(logdir_path)==False:
os.mkdir(logdir_path)
# Load Json & Append Root Path
with open(json_Path, 'r') as json_f:
json_Data = json.load(json_f)
train_Data = json_Data['training']
val_Data = json_Data['validation']
for idx, each_d in enumerate(train_Data):
train_Data[idx]['image'] = os.path.join(data_Root, train_Data[idx]['image'])
for idx, each_d in enumerate(val_Data):
val_Data[idx]['image'] = os.path.join(data_Root, val_Data[idx]['image'])
print('Total Number of Training Data Samples: {}'.format(len(train_Data)))
print(train_Data)
print('#' * 10)
print('Total Number of Validation Data Samples: {}'.format(len(val_Data)))
print(val_Data)
print('#' * 10)
# Set Determinism
set_determinism(seed=123)
# Define Training Transforms
train_Transforms = Compose(
[
LoadImaged(keys=["image"]),
EnsureChannelFirstd(keys=["image"]),
Spacingd(keys=["image"], pixdim=(
2.0, 2.0, 2.0), mode=("bilinear")),
ScaleIntensityRanged(
keys=["image"], a_min=-57, a_max=164,
b_min=0.0, b_max=1.0, clip=True,
),
CropForegroundd(keys=["image"], source_key="image"),
SpatialPadd(keys=["image"], spatial_size=(96, 96, 96)),
RandSpatialCropSamplesd(keys=["image"], roi_size=(96, 96, 96), random_size=False, num_samples=2),
CopyItemsd(keys=["image"], times=2, names=["gt_image", "image_2"], allow_missing_keys=False),
OneOf(transforms=[
RandCoarseDropoutd(keys=["image"], prob=1.0, holes=6, spatial_size=5, dropout_holes=True,
max_spatial_size=32),
RandCoarseDropoutd(keys=["image"], prob=1.0, holes=6, spatial_size=20, dropout_holes=False,
max_spatial_size=64),
]
),
RandCoarseShuffled(keys=["image"], prob=0.8, holes=10, spatial_size=8),
# Please note that that if image, image_2 are called via the same transform call because of the determinism
# they will get augmented the exact same way which is not the required case here, hence two calls are made
OneOf(transforms=[
RandCoarseDropoutd(keys=["image_2"], prob=1.0, holes=6, spatial_size=5, dropout_holes=True,
max_spatial_size=32),
RandCoarseDropoutd(keys=["image_2"], prob=1.0, holes=6, spatial_size=20, dropout_holes=False,
max_spatial_size=64),
]
),
RandCoarseShuffled(keys=["image_2"], prob=0.8, holes=10, spatial_size=8)
]
)
check_ds = Dataset(data=train_Data, transform=train_Transforms)
check_loader = DataLoader(check_ds, batch_size=1)
check_data = first(check_loader)
image = (check_data["image"][0][0])
print(f"image shape: {image.shape}")
# Define Network ViT backbone & Loss & Optimizer
device = torch.device("cuda:0")
model = ViTAutoEnc(
in_channels=1,
img_size=(96, 96, 96),
patch_size=(16, 16, 16),
pos_embed='conv',
hidden_size=768,
mlp_dim=3072,
)
model = model.to(device)
# Define Hyper-paramters for training loop
max_epochs = 500
val_interval = 2
batch_size = 4
lr = 1e-4
epoch_loss_values = []
step_loss_values = []
epoch_cl_loss_values = []
epoch_recon_loss_values = []
val_loss_values = []
best_val_loss = 1000.0
recon_loss = L1Loss()
contrastive_loss = ContrastiveLoss(batch_size=batch_size*2, temperature=0.05)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
# Define DataLoader using MONAI, CacheDataset needs to be used
train_ds = Dataset(data=train_Data, transform=train_Transforms)
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4)
val_ds = Dataset(data=val_Data, transform=train_Transforms)
val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=True, num_workers=4)
for epoch in range(max_epochs):
print("-" * 10)
print(f"epoch {epoch + 1}/{max_epochs}")
model.train()
epoch_loss = 0
epoch_cl_loss = 0
epoch_recon_loss = 0
step = 0
for batch_data in train_loader:
step += 1
start_time = time.time()
inputs, inputs_2, gt_input = (
batch_data["image"].to(device),
batch_data["image_2"].to(device),
batch_data["gt_image"].to(device),
)
optimizer.zero_grad()
outputs_v1, hidden_v1 = model(inputs)
outputs_v2, hidden_v2 = model(inputs_2)
flat_out_v1 = outputs_v1.flatten(start_dim=1, end_dim=4)
flat_out_v2 = outputs_v2.flatten(start_dim=1, end_dim=4)
r_loss = recon_loss(outputs_v1, gt_input)
cl_loss = contrastive_loss(flat_out_v1, flat_out_v2)
# Adjust the CL loss by Recon Loss
total_loss = r_loss + cl_loss * r_loss
total_loss.backward()
optimizer.step()
epoch_loss += total_loss.item()
step_loss_values.append(total_loss.item())
# CL & Recon Loss Storage of Value
epoch_cl_loss += cl_loss.item()
epoch_recon_loss += r_loss.item()
end_time = time.time()
print(
f"{step}/{len(train_ds) // train_loader.batch_size}, "
f"train_loss: {total_loss.item():.4f}, "
f"time taken: {end_time-start_time}s")
epoch_loss /= step
epoch_cl_loss /= step
epoch_recon_loss /= step
epoch_loss_values.append(epoch_loss)
epoch_cl_loss_values.append(epoch_cl_loss)
epoch_recon_loss_values.append(epoch_recon_loss)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
if epoch % val_interval == 0:
print('Entering Validation for epoch: {}'.format(epoch+1))
total_val_loss = 0
val_step = 0
model.eval()
for val_batch in val_loader:
val_step += 1
start_time = time.time()
inputs, gt_input = (
val_batch["image"].to(device),
val_batch["gt_image"].to(device),
)
print('Input shape: {}'.format(inputs.shape))
outputs, outputs_v2 = model(inputs)
val_loss = recon_loss(outputs, gt_input)
total_val_loss += val_loss.item()
end_time = time.time()
total_val_loss /= val_step
val_loss_values.append(total_val_loss)
print(f"epoch {epoch + 1} Validation average loss: {total_val_loss:.4f}, " f"time taken: {end_time-start_time}s")
if total_val_loss < best_val_loss:
print(f"Saving new model based on validation loss {total_val_loss:.4f}")
best_val_loss = total_val_loss
checkpoint = {'epoch': max_epochs,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(checkpoint, os.path.join(logdir_path, 'best_model.pt'))
plt.figure(1, figsize=(8, 8))
plt.subplot(2, 2, 1)
plt.plot(epoch_loss_values)
plt.grid()
plt.title('Training Loss')
plt.subplot(2, 2, 2)
plt.plot(val_loss_values)
plt.grid()
plt.title('Validation Loss')
plt.subplot(2, 2, 3)
plt.plot(epoch_cl_loss_values)
plt.grid()
plt.title('Training Contrastive Loss')
plt.subplot(2, 2, 4)
plt.plot(epoch_recon_loss_values)
plt.grid()
plt.title('Training Recon Loss')
plt.savefig(os.path.join(logdir_path, 'loss_plots.png'))
plt.close(1)
print('Done')
return None
def main(config):
now = datetime.now().strftime("%Y%m%d-%H:%M:%S")
# path
csv_path = config['path']['csv_path']
trained_model_path = config['path'][
'trained_model_path'] # if None, trained from scratch
training_model_folder = os.path.join(
config['path']['training_model_folder'], now) # '/path/to/folder'
if not os.path.exists(training_model_folder):
os.makedirs(training_model_folder)
logdir = os.path.join(training_model_folder, 'logs')
if not os.path.exists(logdir):
os.makedirs(logdir)
# PET CT scan params
image_shape = tuple(config['preprocessing']['image_shape']) # (x, y, z)
in_channels = config['preprocessing']['in_channels']
voxel_spacing = tuple(
config['preprocessing']
['voxel_spacing']) # (4.8, 4.8, 4.8) # in millimeter, (x, y, z)
data_augment = config['preprocessing'][
'data_augment'] # True # for training dataset only
resize = config['preprocessing']['resize'] # True # not use yet
origin = config['preprocessing']['origin'] # how to set the new origin
normalize = config['preprocessing'][
'normalize'] # True # whether or not to normalize the inputs
number_class = config['preprocessing']['number_class'] # 2
# CNN params
architecture = config['model']['architecture'] # 'unet' or 'vnet'
cnn_params = config['model'][architecture]['cnn_params']
# transform list to tuple
for key, value in cnn_params.items():
if isinstance(value, list):
cnn_params[key] = tuple(value)
# Training params
epochs = config['training']['epochs']
batch_size = config['training']['batch_size']
shuffle = config['training']['shuffle']
opt_params = config['training']["optimizer"]["opt_params"]
# Get Data
DM = DataManager(csv_path=csv_path)
train_images_paths, val_images_paths, test_images_paths = DM.get_train_val_test(
wrap_with_dict=True)
# Input preprocessing
# use data augmentation for training
train_transforms = Compose([ # read img + meta info
LoadNifti(keys=["pet_img", "ct_img", "mask_img"]),
Roi2Mask(keys=['pet_img', 'mask_img'],
method='otsu',
tval=0.0,
idx_channel=0),
ResampleReshapeAlign(target_shape=image_shape,
target_voxel_spacing=voxel_spacing,
keys=['pet_img', "ct_img", 'mask_img'],
origin='head',
origin_key='pet_img'),
Sitk2Numpy(keys=['pet_img', 'ct_img', 'mask_img']),
# user can also add other random transforms
RandAffined(keys=("pet_img", "ct_img", "mask_img"),
spatial_size=None,
prob=0.4,
rotate_range=(0, np.pi / 30, np.pi / 15),
shear_range=None,
translate_range=(10, 10, 10),
scale_range=(0.1, 0.1, 0.1),
mode=("bilinear", "bilinear", "nearest"),
padding_mode="border"),
# normalize input
ScaleIntensityRanged(
keys=["pet_img"],
a_min=0.0,
a_max=25.0,
b_min=0.0,
b_max=1.0,
clip=True,
),
ScaleIntensityRanged(
keys=["ct_img"],
a_min=-1000.0,
a_max=1000.0,
b_min=0.0,
b_max=1.0,
clip=True,
),
# Prepare for neural network
ConcatModality(keys=['pet_img', 'ct_img']),
AddChanneld(keys=["mask_img"]), # Add channel to the first axis
ToTensord(keys=["image", "mask_img"]),
])
# without data augmentation for validation
val_transforms = Compose([ # read img + meta info
LoadNifti(keys=["pet_img", "ct_img", "mask_img"]),
Roi2Mask(keys=['pet_img', 'mask_img'],
method='otsu',
tval=0.0,
idx_channel=0),
ResampleReshapeAlign(target_shape=image_shape,
target_voxel_spacing=voxel_spacing,
keys=['pet_img', "ct_img", 'mask_img'],
origin='head',
origin_key='pet_img'),
Sitk2Numpy(keys=['pet_img', 'ct_img', 'mask_img']),
# normalize input
ScaleIntensityRanged(
keys=["pet_img"],
a_min=0.0,
a_max=25.0,
b_min=0.0,
b_max=1.0,
clip=True,
),
ScaleIntensityRanged(
keys=["ct_img"],
a_min=-1000.0,
a_max=1000.0,
b_min=0.0,
b_max=1.0,
clip=True,
),
# Prepare for neural network
ConcatModality(keys=['pet_img', 'ct_img']),
AddChanneld(keys=["mask_img"]), # Add channel to the first axis
ToTensord(keys=["image", "mask_img"]),
])
# create a training data loader
train_ds = monai.data.CacheDataset(data=train_images_paths,
transform=train_transforms,
cache_rate=0.5)
# use batch_size=2 to load images to generate 2 x 4 images for network training
train_loader = monai.data.DataLoader(train_ds,
batch_size=batch_size,
shuffle=shuffle,
num_workers=2)
# create a validation data loader
val_ds = monai.data.CacheDataset(data=val_images_paths,
transform=val_transforms,
cache_rate=1.0)
val_loader = monai.data.DataLoader(val_ds,
batch_size=batch_size,
num_workers=2)
# Model
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = UNet(
dimensions=3, # 3D
in_channels=in_channels,
out_channels=1,
kernel_size=5,
channels=(8, 16, 32, 64, 128),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
loss = monai.losses.DiceLoss(sigmoid=True, squared_pred=True)
opt = torch.optim.Adam(net.parameters(), 1e-3)
# training
val_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
])
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,
"opt": opt
},
save_key_metric=True),
]
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=SimpleInferer(),
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_precision":
Precision(output_transform=lambda x: (x["pred"], x["label"])),
"val_recall":
Recall(output_transform=lambda x: (x["pred"], x["label"]))
},
val_handlers=val_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP evaluation
# amp=True if monai.config.get_torch_version_tuple() >= (1, 6) else False,
)
train_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
])
train_handlers = [
# LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
ValidationHandler(validator=evaluator, interval=1, 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,
prepare_batch=lambda x: (x['image'], x['mask_img']),
inferer=SimpleInferer(),
post_transform=train_post_transforms,
key_train_metric={
"train_mean_dice":
MeanDice(include_background=True,
output_transform=lambda x: (x["pred"], x["label"]))
},
additional_metrics={
"train_acc":
Accuracy(output_transform=lambda x: (x["pred"], x["label"])),
"train_precision":
Precision(output_transform=lambda x: (x["pred"], x["label"])),
"train_recall":
Recall(output_transform=lambda x: (x["pred"], x["label"]))
},
train_handlers=train_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP training
amp=True if monai.config.get_torch_version_tuple() >=
(1, 6) else False,
)
trainer.run()
The image centers of negative samples must be in valid body area.
1. `RandAffined` efficiently performs `rotate`, `scale`, `shear`, `translate`, etc. together based on PyTorch affine transform.
1. `ToTensord` converts the numpy array to PyTorch Tensor for further steps.
"""
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=-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(
Label 1: Bladder
Label 2: Liver
Label 3: Lungs
Label 4: Heart
Label 5: Pancreas
'''
test_transforms = Compose([
LoadImaged(keys="image"),
AddChanneld(keys="image"),
Spacingd(keys="image", pixdim=(1.5, 1.5, 2.0), mode="bilinear"),
Orientationd(keys="image", axcodes="RAS"),
ScaleIntensityRanged(
keys="image",
a_min=-1000,
a_max=300,
b_min=0.0,
b_max=1.0,
clip=True,
),
#CropForegroundd(keys=["image", "label"], source_key="image"),
ToTensord(keys="image"),
])
test_ds = CacheDataset(data=test_files,
transform=test_transforms,
cache_rate=1.0,
num_workers=2)
#test_ds = Dataset(data=test_files)
test_loader = DataLoader(test_ds, batch_size=1, num_workers=2)
"""## Create Model, Loss, Optimizer"""