def train_pre_transforms(self, context: Context):
# Dataset preparation
t: List[Any] = [
LoadImaged(keys=("image", "label")),
AddChanneld(keys=("image", "label")),
SpatialCropForegroundd(keys=("image", "label"),
source_key="label",
spatial_size=self.roi_size),
Resized(keys=("image", "label"),
spatial_size=self.model_size,
mode=("area", "nearest")),
NormalizeIntensityd(keys="image", subtrahend=208.0,
divisor=388.0), # type: ignore
]
if self.dimension == 3:
t.append(FindAllValidSlicesd(label="label", sids="sids"))
t.extend([
AddInitialSeedPointd(label="label",
guidance="guidance",
sids="sids"),
AddGuidanceSignald(image="image", guidance="guidance"),
EnsureTyped(keys=("image", "label"), device=context.device),
SelectItemsd(keys=("image", "label", "guidance")),
])
return t
def pre_transforms(self, data=None) -> Sequence[Callable]:
t = [
LoadImaged(keys="image"),
AsChannelFirstd(keys="image"),
Spacingd(keys="image",
pixdim=[1.0] * self.dimension,
mode="bilinear"),
AddGuidanceFromPointsd(ref_image="image",
guidance="guidance",
dimensions=self.dimension),
]
if self.dimension == 2:
t.append(Fetch2DSliced(keys="image", guidance="guidance"))
t.extend([
AddChanneld(keys="image"),
SpatialCropGuidanced(keys="image",
guidance="guidance",
spatial_size=self.spatial_size),
Resized(keys="image", spatial_size=self.model_size, mode="area"),
ResizeGuidanced(guidance="guidance", ref_image="image"),
NormalizeIntensityd(keys="image", subtrahend=208,
divisor=388), # type: ignore
AddGuidanceSignald(image="image", guidance="guidance"),
EnsureTyped(keys="image",
device=data.get("device") if data else None),
])
return t
def __init__(self, tranforms):
self.tranform_list = []
for tranform in tranforms:
if 'LoadImaged' == tranform:
self.tranform_list.append(LoadImaged(keys=["image", "label"]))
elif 'AsChannelFirstd' == tranform:
self.tranform_list.append(AsChannelFirstd(keys="image"))
elif 'ConvertToMultiChannelBasedOnBratsClassesd' == tranform:
self.tranform_list.append(ConvertToMultiChannelBasedOnBratsClassesd(keys="label"))
elif 'Spacingd' == tranform:
self.tranform_list.append(Spacingd(keys=["image", "label"], pixdim=(1.5, 1.5, 2.0), mode=("bilinear", "nearest")))
elif 'Orientationd' == tranform:
self.tranform_list.append(Orientationd(keys=["image", "label"], axcodes="RAS"))
elif 'CenterSpatialCropd' == tranform:
self.tranform_list.append(CenterSpatialCropd(keys=["image", "label"], roi_size=[128, 128, 64]))
elif 'NormalizeIntensityd' == tranform:
self.tranform_list.append(NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True))
elif 'ToTensord' == tranform:
self.tranform_list.append(ToTensord(keys=["image", "label"]))
elif 'Activations' == tranform:
self.tranform_list.append(Activations(sigmoid=True))
elif 'AsDiscrete' == tranform:
self.tranform_list.append(AsDiscrete(threshold_values=True))
else:
raise ValueError(
f"Unsupported tranform: {tranform}. Please add it to support it."
)
super().__init__(self.tranform_list)
def test_channel_wise(self):
key = 'img'
normalizer = NormalizeIntensityd(keys=key,
nonzero=True,
channel_wise=True)
input_data = {key: np.array([[0., 3., 0., 4.], [0., 4., 0., 5.]])}
expected = np.array([[0., -1., 0., 1.], [0., -1., 0., 1.]])
np.testing.assert_allclose(expected, normalizer(input_data)[key])
def test_nonzero(self, input_param, input_data, expected_data):
key = "img"
normalizer = NormalizeIntensityd(**input_param)
normalized = normalizer(input_data)[key]
self.assertEqual(type(input_data[key]), type(normalized))
if isinstance(normalized, torch.Tensor):
self.assertEqual(input_data[key].device, normalized.device)
assert_allclose(normalized, expected_data)
def test_channel_wise(self):
key = "img"
normalizer = NormalizeIntensityd(keys=key,
nonzero=True,
channel_wise=True)
input_data = {
key: np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0, 5.0]])
}
expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]])
np.testing.assert_allclose(expected, normalizer(input_data)[key])
def get_transforms(self):
self.logger.info("Getting transforms...")
# Setup transforms of data sets
train_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
NormalizeIntensityd(keys=["image"]),
SpatialPadd(keys=["image", "label"],
spatial_size=self.pad_crop_shape),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandSpatialCropd(keys=["image", "label"],
roi_size=self.pad_crop_shape,
random_center=True,
random_size=False),
ToTensord(keys=["image", "label"]),
])
val_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
NormalizeIntensityd(keys=["image"]),
SpatialPadd(keys=["image", "label"],
spatial_size=self.pad_crop_shape),
RandSpatialCropd(
keys=["image", "label"],
roi_size=self.pad_crop_shape,
random_center=True,
random_size=False,
),
ToTensord(keys=["image", "label"]),
])
test_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
NormalizeIntensityd(keys=["image"]),
ToTensord(keys=["image", "label"]),
])
return train_transforms, val_transforms, test_transforms
def test_image_normalize_intensityd(self, im_type):
key = "img"
im = im_type(self.imt)
normalizer = NormalizeIntensityd(keys=[key])
normalized = normalizer({key: im})[key]
expected = (self.imt - np.mean(self.imt)) / np.std(self.imt)
self.assertEqual(type(im), type(normalized))
if isinstance(normalized, torch.Tensor):
self.assertEqual(im.device, normalized.device)
assert_allclose(normalized, im_type(expected), rtol=1e-3)
def test_image_normalize_intensityd(self, im_type):
key = "img"
im = im_type(self.imt)
normalizer = NormalizeIntensityd(keys=[key])
normalized = normalizer({key: im})[key]
expected = (self.imt - np.mean(self.imt)) / np.std(self.imt)
assert_allclose(normalized,
im_type(expected),
rtol=1e-3,
type_test="tensor")
def test_channel_wise(self, im_type):
key = "img"
normalizer = NormalizeIntensityd(keys=key,
nonzero=True,
channel_wise=True)
input_data = {
key: im_type(np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0,
5.0]]))
}
normalized = normalizer(input_data)[key]
expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]])
assert_allclose(normalized, im_type(expected), type_test="tensor")
def get_test_transforms(end_image_shape):
test_transforms = Compose([
LoadNiftid(keys=["image"]),
AddChanneld(keys=["image"]),
Orientationd(keys=["image"], axcodes="RAS"),
Winsorized(keys=["image"]),
NormalizeIntensityd(keys=["image"]),
ScaleIntensityd(keys=["image"]),
SpatialPadd(keys=["image"], spatial_size=end_image_shape),
ToTensord(keys=["image"]),
])
return test_transforms
def pre_transforms(self, data):
return [
LoadImaged(keys="image"),
AsChannelFirstd(keys="image"),
Spacingd(keys="image", pixdim=[1.0, 1.0, 1.0], mode="bilinear"),
AddGuidanceFromPointsd(ref_image="image", guidance="guidance", dimensions=3),
AddChanneld(keys="image"),
SpatialCropGuidanced(keys="image", guidance="guidance", spatial_size=self.spatial_size),
Resized(keys="image", spatial_size=self.model_size, mode="area"),
ResizeGuidanced(guidance="guidance", ref_image="image"),
NormalizeIntensityd(keys="image", subtrahend=208, divisor=388),
AddGuidanceSignald(image="image", guidance="guidance"),
]
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 test_channel_wise(self, im_type):
key = "img"
normalizer = NormalizeIntensityd(keys=key,
nonzero=True,
channel_wise=True)
input_data = {
key: im_type(np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0,
5.0]]))
}
normalized = normalizer(input_data)[key]
self.assertEqual(type(input_data[key]), type(normalized))
if isinstance(normalized, torch.Tensor):
self.assertEqual(input_data[key].device, normalized.device)
expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]])
assert_allclose(normalized, im_type(expected))
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 pre_transforms(self, data=None):
return [
LoadImaged(keys="image"),
AsChannelFirstd(keys="image"),
Spacingd(keys="image", pixdim=[1.0, 1.0], mode="bilinear"),
AddGuidanceFromPointsd(ref_image="image",
guidance="guidance",
dimensions=2),
Fetch2DSliced(keys="image", guidance="guidance"),
AddChanneld(keys="image"),
SpatialCropGuidanced(keys="image",
guidance="guidance",
spatial_size=[256, 256]),
Resized(keys="image", spatial_size=[256, 256], mode="area"),
ResizeGuidanced(guidance="guidance", ref_image="image"),
NormalizeIntensityd(keys="image", subtrahend=208,
divisor=388), # type: ignore
AddGuidanceSignald(image="image", guidance="guidance"),
]
def get_pre_transforms(roi_size, model_size, dimensions):
t = [
LoadImaged(keys=('image', 'label')),
AddChanneld(keys=('image', 'label')),
SpatialCropForegroundd(keys=('image', 'label'),
source_key='label',
spatial_size=roi_size),
Resized(keys=('image', 'label'),
spatial_size=model_size,
mode=('area', 'nearest')),
NormalizeIntensityd(keys='image', subtrahend=208.0, divisor=388.0)
]
if dimensions == 3:
t.append(FindAllValidSlicesd(label='label', sids='sids'))
t.extend([
AddInitialSeedPointd(label='label', guidance='guidance', sids='sids'),
AddGuidanceSignald(image='image', guidance='guidance'),
ToTensord(keys=('image', 'label'))
])
return Compose(t)
def get_pre_transforms(roi_size, model_size, dimensions):
t = [
LoadImaged(keys=("image", "label")),
AddChanneld(keys=("image", "label")),
SpatialCropForegroundd(keys=("image", "label"),
source_key="label",
spatial_size=roi_size),
Resized(keys=("image", "label"),
spatial_size=model_size,
mode=("area", "nearest")),
NormalizeIntensityd(keys="image", subtrahend=208.0, divisor=388.0),
]
if dimensions == 3:
t.append(FindAllValidSlicesd(label="label", sids="sids"))
t.extend([
AddInitialSeedPointd(label="label", guidance="guidance", sids="sids"),
AddGuidanceSignald(image="image", guidance="guidance"),
EnsureTyped(keys=("image", "label")),
])
return Compose(t)
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
'label': label_name
} for image_name, label_name in zip(train_images, train_segs)]
monai_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
Orientationd(keys=["image", "label"], axcodes="RAS"),
# ThresholdIntensityd(keys=['image'], threshold=-135, above=True, cval=-135),
# ThresholdIntensityd(keys=['image'], threshold=215, above=False, cval=215),
CropForegroundd(
keys=['image', 'label'],
source_key='image',
start_coord_key='foreground_start_coord',
end_coord_key='foreground_end_coord',
), # crop CropForeground
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
# Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')),
SpatialPadd(keys=['image', 'label'],
spatial_size=opt.patch_size,
method='end'),
RandSpatialCropd(keys=['image', 'label'],
roi_size=opt.patch_size,
random_size=False),
ToTensord(keys=[
'image', 'label', 'foreground_start_coord', 'foreground_end_coord'
], )
]
transform = Compose(monai_transforms)
check_ds = monai.data.Dataset(data=data_dicts, transform=transform)
def test_nonzero(self, input_param, input_data, expected_data):
normalizer = NormalizeIntensityd(**input_param)
np.testing.assert_allclose(expected_data,
normalizer(input_data)['img'])
def test_image_normalize_intensityd(self):
key = 'img'
normalizer = NormalizeIntensityd(keys=[key])
normalized = normalizer({key: self.imt})
expected = (self.imt - np.mean(self.imt)) / np.std(self.imt)
np.testing.assert_allclose(normalized[key], expected, rtol=1e-6)
def main():
opt = Options().parse()
# monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
set_determinism(seed=0)
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
device = torch.device(opt.gpu_id)
# ------- Data loader creation ----------
# images
images = sorted(glob(os.path.join(opt.images_folder, 'image*.nii')))
segs = sorted(glob(os.path.join(opt.labels_folder, 'label*.nii')))
train_files = []
val_files = []
for i in range(opt.models_ensemble):
train_files.append([{
"image": img,
"label": seg
} for img, seg in zip(
images[:(opt.split_val * i)] +
images[(opt.split_val *
(i + 1)):(len(images) -
opt.split_val)], segs[:(opt.split_val * i)] +
segs[(opt.split_val * (i + 1)):(len(images) - opt.split_val)])])
val_files.append([{
"image": img,
"label": seg
} for img, seg in zip(
images[(opt.split_val * i):(opt.split_val *
(i + 1))], segs[(opt.split_val *
i):(opt.split_val *
(i + 1))])])
test_files = [{
"image": img,
"label": seg
} for img, seg in zip(images[(len(images) -
opt.split_test):len(images)], segs[(
len(images) -
opt.split_test):len(images)])]
# ----------- Transforms list --------------
if opt.resolution is not None:
train_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['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 = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
Spacingd(keys=['image', 'label'],
pixdim=opt.resolution,
mode=('bilinear', 'nearest')),
ToTensord(keys=['image', 'label'])
]
else:
train_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['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 = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
ToTensord(keys=['image', 'label'])
]
train_transforms = Compose(train_transforms)
val_transforms = Compose(val_transforms)
# ---------- Creation of DataLoaders -------------
train_dss = [
CacheDataset(data=train_files[i], transform=train_transforms)
for i in range(opt.models_ensemble)
]
train_loaders = [
DataLoader(train_dss[i],
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.workers,
pin_memory=torch.cuda.is_available())
for i in range(opt.models_ensemble)
]
val_dss = [
CacheDataset(data=val_files[i], transform=val_transforms)
for i in range(opt.models_ensemble)
]
val_loaders = [
DataLoader(val_dss[i],
batch_size=1,
num_workers=opt.workers,
pin_memory=torch.cuda.is_available())
for i in range(opt.models_ensemble)
]
test_ds = CacheDataset(data=test_files, transform=val_transforms)
test_loader = DataLoader(test_ds,
batch_size=1,
num_workers=opt.workers,
pin_memory=torch.cuda.is_available())
def train(index):
# ---------- Build the nn-Unet network ------------
if opt.resolution is None:
sizes, spacings = opt.patch_size, opt.spacing
else:
sizes, spacings = opt.patch_size, opt.resolution
strides, kernels = [], []
while True:
spacing_ratio = [sp / min(spacings) for sp in spacings]
stride = [
2 if ratio <= 2 and size >= 8 else 1
for (ratio, size) in zip(spacing_ratio, sizes)
]
kernel = [3 if ratio <= 2 else 1 for ratio in spacing_ratio]
if all(s == 1 for s in stride):
break
sizes = [i / j for i, j in zip(sizes, stride)]
spacings = [i * j for i, j in zip(spacings, stride)]
kernels.append(kernel)
strides.append(stride)
strides.insert(0, len(spacings) * [1])
kernels.append(len(spacings) * [3])
net = monai.networks.nets.DynUNet(
spatial_dims=3,
in_channels=opt.in_channels,
out_channels=opt.out_channels,
kernel_size=kernels,
strides=strides,
upsample_kernel_size=strides[1:],
res_block=True,
# act=act_type,
# norm=Norm.BATCH,
).to(device)
from torch.autograd import Variable
from torchsummaryX import summary
data = Variable(
torch.randn(int(opt.batch_size), int(opt.in_channels),
int(opt.patch_size[0]), int(opt.patch_size[1]),
int(opt.patch_size[2]))).cuda()
out = net(data)
summary(net, data)
print("out size: {}".format(out.size()))
# if opt.preload is not None:
# net.load_state_dict(torch.load(opt.preload))
# ---------- ------------------------ ------------
optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
optim, lr_lambda=lambda epoch: (1 - epoch / opt.epochs)**0.9)
loss_function = monai.losses.DiceCELoss(sigmoid=True)
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),
CheckpointSaver(save_dir="./runs/",
save_dict={"net": net},
save_key_metric=True),
]
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loaders[index],
network=net,
inferer=SlidingWindowInferer(roi_size=opt.patch_size,
sw_batch_size=opt.batch_size,
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"]),
)
},
val_handlers=val_handlers)
train_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
# KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
])
train_handlers = [
ValidationHandler(validator=evaluator,
interval=5,
epoch_level=True),
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir="./runs/",
save_dict={
"net": net,
"opt": optim
},
save_final=True,
epoch_level=True),
]
trainer = SupervisedTrainer(
device=device,
max_epochs=opt.epochs,
train_data_loader=train_loaders[index],
network=net,
optimizer=optim,
loss_function=loss_function,
inferer=SimpleInferer(),
post_transform=train_post_transforms,
amp=False,
train_handlers=train_handlers,
)
trainer.run()
return net
models = [train(i) for i in range(opt.models_ensemble)]
# -------- Test the models ---------
def ensemble_evaluate(post_transforms, models):
evaluator = EnsembleEvaluator(
device=device,
val_data_loader=test_loader,
pred_keys=opt.pred_keys,
networks=models,
inferer=SlidingWindowInferer(roi_size=opt.patch_size,
sw_batch_size=opt.batch_size,
overlap=0.5),
post_transform=post_transforms,
key_val_metric={
"test_mean_dice":
MeanDice(
include_background=True,
output_transform=lambda x: (x["pred"], x["label"]),
)
},
)
evaluator.run()
mean_post_transforms = Compose([
MeanEnsembled(
keys=opt.pred_keys,
output_key="pred",
# in this particular example, we use validation metrics as weights
weights=opt.weights_models,
),
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
# KeepLargestConnectedComponentd(keys="pred", applied_labels=[1])
])
print('Results from MeanEnsembled:')
ensemble_evaluate(mean_post_transforms, models)
vote_post_transforms = Compose([
Activationsd(keys=opt.pred_keys, sigmoid=True),
# transform data into discrete before voting
AsDiscreted(keys=opt.pred_keys, threshold_values=True),
# KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
VoteEnsembled(keys=opt.pred_keys, output_key="pred"),
])
print('Results from VoteEnsembled:')
ensemble_evaluate(vote_post_transforms, models)
def run_training(train_file_list, valid_file_list, config_info):
"""
Pipeline to train a dynUNet segmentation model in MONAI. It is composed of the following main blocks:
* Data Preparation: Extract the filenames and prepare the training/validation processing transforms
* Load Data: Load training and validation data to PyTorch DataLoader
* Network Preparation: Define the network, loss function, optimiser and learning rate scheduler
* MONAI Evaluator: Initialise the dynUNet evaluator, i.e. the class providing utilities to perform validation
during training. Attach handlers to save the best model on the validation set. A 2D sliding window approach
on the 3D volume is used at evaluation. The mean 3D Dice is used as validation metric.
* MONAI Trainer: Initialise the dynUNet trainer, i.e. the class providing utilities to perform the training loop.
* Run training: The MONAI trainer is run, performing training and validation during training.
Args:
train_file_list: .txt or .csv file (with no header) storing two-columns filenames for training:
image filename in the first column and segmentation filename in the second column.
The two columns should be separated by a comma.
See monaifbs/config/mock_train_file_list_for_dynUnet_training.txt for an example of the expected format.
valid_file_list: .txt or .csv file (with no header) storing two-columns filenames for validation:
image filename in the first column and segmentation filename in the second column.
The two columns should be separated by a comma.
See monaifbs/config/mock_valid_file_list_for_dynUnet_training.txt for an example of the expected format.
config_info: dict, contains configuration parameters for sampling, network and training.
See monaifbs/config/monai_dynUnet_training_config.yml for an example of the expected fields.
"""
"""
Read input and configuration parameters
"""
# print MONAI config information
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print_config()
# print to log the parameter setups
print(yaml.dump(config_info))
# extract network parameters, perform checks/set defaults if not present and print them to log
if 'seg_labels' in config_info['training'].keys():
seg_labels = config_info['training']['seg_labels']
else:
seg_labels = [1]
nr_out_channels = len(seg_labels)
print("Considering the following {} labels in the segmentation: {}".format(nr_out_channels, seg_labels))
patch_size = config_info["training"]["inplane_size"] + [1]
print("Considering patch size = {}".format(patch_size))
spacing = config_info["training"]["spacing"]
print("Bringing all images to spacing = {}".format(spacing))
if 'model_to_load' in config_info['training'].keys() and config_info['training']['model_to_load'] is not None:
model_to_load = config_info['training']['model_to_load']
if not os.path.exists(model_to_load):
raise FileNotFoundError("Cannot find model: {}".format(model_to_load))
else:
print("Loading model from {}".format(model_to_load))
else:
model_to_load = None
# set up either GPU or CPU usage
if torch.cuda.is_available():
print("\n#### GPU INFORMATION ###")
print("Using device number: {}, name: {}\n".format(torch.cuda.current_device(), torch.cuda.get_device_name()))
current_device = torch.device("cuda:0")
else:
current_device = torch.device("cpu")
print("Using device: {}".format(current_device))
# set determinism if required
if 'manual_seed' in config_info['training'].keys() and config_info['training']['manual_seed'] is not None:
seed = config_info['training']['manual_seed']
else:
seed = None
if seed is not None:
print("Using determinism with seed = {}\n".format(seed))
set_determinism(seed=seed)
"""
Setup data output directory
"""
out_model_dir = os.path.join(config_info['output']['out_dir'],
datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + '_' +
config_info['output']['out_postfix'])
print("Saving to directory {}\n".format(out_model_dir))
# create cache directory to store results for Persistent Dataset
if 'cache_dir' in config_info['output'].keys():
out_cache_dir = config_info['output']['cache_dir']
else:
out_cache_dir = os.path.join(out_model_dir, 'persistent_cache')
persistent_cache: Path = Path(out_cache_dir)
persistent_cache.mkdir(parents=True, exist_ok=True)
"""
Data preparation
"""
# Read the input files for training and validation
print("*** Loading input data for training...")
train_files = create_data_list_of_dictionaries(train_file_list)
print("Number of inputs for training = {}".format(len(train_files)))
val_files = create_data_list_of_dictionaries(valid_file_list)
print("Number of inputs for validation = {}".format(len(val_files)))
# Define MONAI processing transforms for the training data. This includes:
# - Load Nifti files and convert to format Batch x Channel x Dim1 x Dim2 x Dim3
# - CropForegroundd: Reduce the background from the MR image
# - InPlaneSpacingd: Perform in-plane resampling to the desired spacing, but preserve the resolution along the
# last direction (lowest resolution) to avoid introducing motion artefact resampling errors
# - SpatialPadd: Pad the in-plane size to the defined network input patch size [N, M] if needed
# - NormalizeIntensityd: Apply whitening
# - RandSpatialCropd: Crop a random patch from the input with size [B, C, N, M, 1]
# - SqueezeDimd: Convert the 3D patch to a 2D one as input to the network (i.e. bring it to size [B, C, N, M])
# - Apply data augmentation (RandZoomd, RandRotated, RandGaussianNoised, RandGaussianSmoothd, RandScaleIntensityd,
# RandFlipd)
# - ToTensor: convert to pytorch tensor
train_transforms = Compose(
[
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
CropForegroundd(keys=["image", "label"], source_key="image"),
InPlaneSpacingd(
keys=["image", "label"],
pixdim=spacing,
mode=("bilinear", "nearest"),
),
SpatialPadd(keys=["image", "label"], spatial_size=patch_size,
mode=["constant", "edge"]),
NormalizeIntensityd(keys=["image"], nonzero=False, channel_wise=True),
RandSpatialCropd(keys=["image", "label"], roi_size=patch_size, random_size=False),
SqueezeDimd(keys=["image", "label"], dim=-1),
RandZoomd(
keys=["image", "label"],
min_zoom=0.9,
max_zoom=1.2,
mode=("bilinear", "nearest"),
align_corners=(True, None),
prob=0.16,
),
RandRotated(keys=["image", "label"], range_x=90, range_y=90, prob=0.2,
keep_size=True, mode=["bilinear", "nearest"],
padding_mode=["zeros", "border"]),
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, 1], prob=0.5),
ToTensord(keys=["image", "label"]),
]
)
# Define MONAI processing transforms for the validation data
# - Load Nifti files and convert to format Batch x Channel x Dim1 x Dim2 x Dim3
# - CropForegroundd: Reduce the background from the MR image
# - InPlaneSpacingd: Perform in-plane resampling to the desired spacing, but preserve the resolution along the
# last direction (lowest resolution) to avoid introducing motion artefact resampling errors
# - SpatialPadd: Pad the in-plane size to the defined network input patch size [N, M] if needed
# - NormalizeIntensityd: Apply whitening
# - ToTensor: convert to pytorch tensor
# NOTE: The validation data is kept 3D as a 2D sliding window approach is used throughout the volume at inference
val_transforms = Compose(
[
LoadNiftid(keys=["image", "label"]),
AddChanneld(keys=["image", "label"]),
CropForegroundd(keys=["image", "label"], source_key="image"),
InPlaneSpacingd(
keys=["image", "label"],
pixdim=spacing,
mode=("bilinear", "nearest"),
),
SpatialPadd(keys=["image", "label"], spatial_size=patch_size, mode=["constant", "edge"]),
NormalizeIntensityd(keys=["image"], nonzero=False, channel_wise=True),
ToTensord(keys=["image", "label"]),
]
)
"""
Load data
"""
# create training data loader
train_ds = PersistentDataset(data=train_files, transform=train_transforms,
cache_dir=persistent_cache)
train_loader = DataLoader(train_ds,
batch_size=config_info['training']['batch_size_train'],
shuffle=True,
num_workers=config_info['device']['num_workers'])
check_train_data = misc.first(train_loader)
print("Training data tensor shapes:")
print("Image = {}; Label = {}".format(check_train_data["image"].shape, check_train_data["label"].shape))
# create validation data loader
if config_info['training']['batch_size_valid'] != 1:
raise Exception("Batch size different from 1 at validation ar currently not supported")
val_ds = PersistentDataset(data=val_files, transform=val_transforms, cache_dir=persistent_cache)
val_loader = DataLoader(val_ds,
batch_size=1,
shuffle=False,
num_workers=config_info['device']['num_workers'])
check_valid_data = misc.first(val_loader)
print("Validation data tensor shapes (Example):")
print("Image = {}; Label = {}\n".format(check_valid_data["image"].shape, check_valid_data["label"].shape))
"""
Network preparation
"""
print("*** Preparing the network ...")
# automatically extracts the strides and kernels based on nnU-Net empirical rules
spacings = spacing[:2]
sizes = patch_size[:2]
strides, kernels = [], []
while True:
spacing_ratio = [sp / min(spacings) for sp in spacings]
stride = [2 if ratio <= 2 and size >= 8 else 1 for (ratio, size) in zip(spacing_ratio, sizes)]
kernel = [3 if ratio <= 2 else 1 for ratio in spacing_ratio]
if all(s == 1 for s in stride):
break
sizes = [i / j for i, j in zip(sizes, stride)]
spacings = [i * j for i, j in zip(spacings, stride)]
kernels.append(kernel)
strides.append(stride)
strides.insert(0, len(spacings) * [1])
kernels.append(len(spacings) * [3])
# initialise the network
net = DynUNet(
spatial_dims=2,
in_channels=1,
out_channels=nr_out_channels,
kernel_size=kernels,
strides=strides,
upsample_kernel_size=strides[1:],
norm_name="instance",
deep_supervision=True,
deep_supr_num=2,
res_block=False,
).to(current_device)
print(net)
# define the loss function
loss_function = choose_loss_function(nr_out_channels, config_info)
# define the optimiser and the learning rate scheduler
opt = torch.optim.SGD(net.parameters(), lr=float(config_info['training']['lr']), momentum=0.95)
scheduler = torch.optim.lr_scheduler.LambdaLR(
opt, lr_lambda=lambda epoch: (1 - epoch / config_info['training']['nr_train_epochs']) ** 0.9
)
"""
MONAI evaluator
"""
print("*** Preparing the dynUNet evaluator engine...\n")
# val_post_transforms = Compose(
# [
# Activationsd(keys="pred", sigmoid=True),
# ]
# )
val_handlers = [
StatsHandler(output_transform=lambda x: None),
TensorBoardStatsHandler(log_dir=os.path.join(out_model_dir, "valid"),
output_transform=lambda x: None,
global_epoch_transform=lambda x: trainer.state.iteration),
CheckpointSaver(save_dir=out_model_dir, save_dict={"net": net, "opt": opt}, save_key_metric=True,
file_prefix='best_valid'),
]
if config_info['output']['val_image_to_tensorboad']:
val_handlers.append(TensorBoardImageHandler(log_dir=os.path.join(out_model_dir, "valid"),
batch_transform=lambda x: (x["image"], x["label"]),
output_transform=lambda x: x["pred"], interval=2))
# Define customized evaluator
class DynUNetEvaluator(SupervisedEvaluator):
def _iteration(self, engine, batchdata):
inputs, targets = self.prepare_batch(batchdata)
inputs, targets = inputs.to(engine.state.device), targets.to(engine.state.device)
flip_inputs_1 = torch.flip(inputs, dims=(2,))
flip_inputs_2 = torch.flip(inputs, dims=(3,))
flip_inputs_3 = torch.flip(inputs, dims=(2, 3))
def _compute_pred():
pred = self.inferer(inputs, self.network)
# use random flipping as data augmentation at inference
flip_pred_1 = torch.flip(self.inferer(flip_inputs_1, self.network), dims=(2,))
flip_pred_2 = torch.flip(self.inferer(flip_inputs_2, self.network), dims=(3,))
flip_pred_3 = torch.flip(self.inferer(flip_inputs_3, self.network), dims=(2, 3))
return (pred + flip_pred_1 + flip_pred_2 + flip_pred_3) / 4
# execute forward computation
self.network.eval()
with torch.no_grad():
if self.amp:
with torch.cuda.amp.autocast():
predictions = _compute_pred()
else:
predictions = _compute_pred()
return {"image": inputs, "label": targets, "pred": predictions}
evaluator = DynUNetEvaluator(
device=current_device,
val_data_loader=val_loader,
network=net,
inferer=SlidingWindowInferer2D(roi_size=patch_size, sw_batch_size=4, overlap=0.0),
post_transform=None,
key_val_metric={
"Mean_dice": MeanDice(
include_background=False,
to_onehot_y=True,
mutually_exclusive=True,
output_transform=lambda x: (x["pred"], x["label"]),
)
},
val_handlers=val_handlers,
amp=False,
)
"""
MONAI trainer
"""
print("*** Preparing the dynUNet trainer engine...\n")
# train_post_transforms = Compose(
# [
# Activationsd(keys="pred", sigmoid=True),
# ]
# )
validation_every_n_epochs = config_info['training']['validation_every_n_epochs']
epoch_len = len(train_ds) // train_loader.batch_size
validation_every_n_iters = validation_every_n_epochs * epoch_len
# define event handlers for the trainer
writer_train = SummaryWriter(log_dir=os.path.join(out_model_dir, "train"))
train_handlers = [
LrScheduleHandler(lr_scheduler=scheduler, print_lr=True),
ValidationHandler(validator=evaluator, interval=validation_every_n_iters, epoch_level=False),
StatsHandler(tag_name="train_loss", output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(summary_writer=writer_train,
log_dir=os.path.join(out_model_dir, "train"), tag_name="Loss",
output_transform=lambda x: x["loss"],
global_epoch_transform=lambda x: trainer.state.iteration),
CheckpointSaver(save_dir=out_model_dir, save_dict={"net": net, "opt": opt},
save_final=True,
save_interval=2, epoch_level=True,
n_saved=config_info['output']['max_nr_models_saved']),
]
if model_to_load is not None:
train_handlers.append(CheckpointLoader(load_path=model_to_load, load_dict={"net": net, "opt": opt}))
# define customized trainer
class DynUNetTrainer(SupervisedTrainer):
def _iteration(self, engine, batchdata):
inputs, targets = self.prepare_batch(batchdata)
inputs, targets = inputs.to(engine.state.device), targets.to(engine.state.device)
def _compute_loss(preds, label):
labels = [label] + [interpolate(label, pred.shape[2:]) for pred in preds[1:]]
return sum([0.5 ** i * self.loss_function(p, l) for i, (p, l) in enumerate(zip(preds, labels))])
self.network.train()
self.optimizer.zero_grad()
if self.amp and self.scaler is not None:
with torch.cuda.amp.autocast():
predictions = self.inferer(inputs, self.network)
loss = _compute_loss(predictions, targets)
self.scaler.scale(loss).backward()
self.scaler.step(self.optimizer)
self.scaler.update()
else:
predictions = self.inferer(inputs, self.network)
loss = _compute_loss(predictions, targets).mean()
loss.backward()
self.optimizer.step()
return {"image": inputs, "label": targets, "pred": predictions, "loss": loss.item()}
trainer = DynUNetTrainer(
device=current_device,
max_epochs=config_info['training']['nr_train_epochs'],
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=loss_function,
inferer=SimpleInferer(),
post_transform=None,
key_train_metric=None,
train_handlers=train_handlers,
amp=False,
)
"""
Run training
"""
print("*** Run training...")
trainer.run()
print("Done!")
def main_worker(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
if not os.path.exists(args.dir):
raise FileNotFoundError(f"Missing directory {args.dir}")
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
total_start = time.time()
train_transforms = Compose([
# load 4 Nifti images and stack them together
LoadNiftid(keys=["image", "label"]),
AsChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Spacingd(keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest")),
Orientationd(keys=["image", "label"], axcodes="RAS"),
RandSpatialCropd(keys=["image", "label"],
roi_size=[128, 128, 64],
random_size=False),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
ToTensord(keys=["image", "label"]),
])
# create a training data loader
train_ds = BratsCacheDataset(
root_dir=args.dir,
transform=train_transforms,
section="training",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=True,
)
train_loader = DataLoader(train_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
# validation transforms and dataset
val_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AsChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Spacingd(keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest")),
Orientationd(keys=["image", "label"], axcodes="RAS"),
CenterSpatialCropd(keys=["image", "label"], roi_size=[128, 128, 64]),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
ToTensord(keys=["image", "label"]),
])
val_ds = BratsCacheDataset(
root_dir=args.dir,
transform=val_transforms,
section="validation",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=False,
)
val_loader = DataLoader(val_ds,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if dist.get_rank() == 0:
# Logging for TensorBoard
writer = SummaryWriter(log_dir=args.log_dir)
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{args.local_rank}")
if args.network == "UNet":
model = UNet(
dimensions=3,
in_channels=4,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
else:
model = SegResNet(in_channels=4,
out_channels=3,
init_filters=16,
dropout_prob=0.2).to(device)
loss_function = DiceLoss(to_onehot_y=False,
sigmoid=True,
squared_pred=True)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=1e-5,
amsgrad=True)
# wrap the model with DistributedDataParallel module
model = DistributedDataParallel(model, device_ids=[args.local_rank])
# start a typical PyTorch training
total_epoch = args.epochs
best_metric = -1000000
best_metric_epoch = -1
epoch_time = AverageMeter("Time", ":6.3f")
progress = ProgressMeter(total_epoch, [epoch_time], prefix="Epoch: ")
end = time.time()
print(f"Time elapsed before training: {end-total_start}")
for epoch in range(total_epoch):
train_loss = train(train_loader, model, loss_function, optimizer,
epoch, args, device)
epoch_time.update(time.time() - end)
if epoch % args.print_freq == 0:
progress.display(epoch)
if dist.get_rank() == 0:
writer.add_scalar("Loss/train", train_loss, epoch)
if (epoch + 1) % args.val_interval == 0:
metric, metric_tc, metric_wt, metric_et = evaluate(
model, val_loader, device)
if dist.get_rank() == 0:
writer.add_scalar("Mean Dice/val", metric, epoch)
writer.add_scalar("Mean Dice TC/val", metric_tc, epoch)
writer.add_scalar("Mean Dice WT/val", metric_wt, epoch)
writer.add_scalar("Mean Dice ET/val", metric_et, epoch)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
print(
f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
f" tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f}"
f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
end = time.time()
print(f"Time elapsed after epoch {epoch + 1} is {end - total_start}")
if dist.get_rank() == 0:
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
# 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")
writer.flush()
dist.destroy_process_group()
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
# Need to concatenate multiple channels here if you want multichannel segmentation
# Check other examples on Monai webpage.
if opt.resolution is not None:
train_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['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 = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['image']),
Spacingd(keys=['image', 'label'], pixdim=opt.resolution, mode=('bilinear', 'nearest')),
ToTensord(keys=['image', 'label'])
]
else:
train_transforms = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['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 = [
LoadImaged(keys=['image', 'label']),
AddChanneld(keys=['image', 'label']),
NormalizeIntensityd(keys=['image']),
ScaleIntensityd(keys=['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, reduction="mean")
post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
# loss_function = monai.losses.DiceLoss(sigmoid=True)
# loss_function = monai.losses.TverskyLoss(sigmoid=True, alpha=0.3, beta=0.7)
loss_function = monai.losses.DiceCELoss(sigmoid=True)
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)
val_outputs = post_trans(val_outputs)
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 run_inference(input_data, config_info):
"""
Pipeline to run inference with MONAI dynUNet model. The pipeline reads the input filenames, applies the required
preprocessing and creates the pytorch dataloader; it then performs evaluation on each input file using a trained
dynUNet model (random flipping augmentation is applied at inference).
It uses the dynUNet model implemented in the MONAI framework
(https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/nets/dynunet.py)
which is inspired by the nnU-Net framework (https://arxiv.org/abs/1809.10486)
Inference is performed in 2D slice-by-slice, all slices are then recombined together into the 3D volume.
Args:
input_data: str or list of strings, filenames of images to be processed
config_info: dict, contains the configuration parameters to reload the trained model
"""
"""
Read input and configuration parameters
"""
val_files = create_data_list_of_dictionaries(input_data)
# print MONAI config information
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
print("*** MONAI config: ")
print_config()
# print to log the parameter setups
print("*** Network inference config: ")
print(yaml.dump(config_info))
# inference params
nr_out_channels = config_info['inference']['nr_out_channels']
spacing = config_info["inference"]["spacing"]
prob_thr = config_info['inference']['probability_threshold']
model_to_load = config_info['inference']['model_to_load']
if not os.path.exists(model_to_load):
raise FileNotFoundError('Trained model not found')
patch_size = config_info["inference"]["inplane_size"] + [1]
print("Considering patch size = {}".format(patch_size))
# set up either GPU or CPU usage
if torch.cuda.is_available():
print("\n#### GPU INFORMATION ###")
print("Using device number: {}, name: {}".format(
torch.cuda.current_device(), torch.cuda.get_device_name()))
current_device = torch.device("cuda:0")
else:
current_device = torch.device("cpu")
print("Using device: {}".format(current_device))
"""
Data Preparation
"""
print("*** Preparing data ... ")
# data preprocessing for inference:
# - convert data to right format [batch, channel, dim, dim, dim]
# - resample to the training resolution in-plane (not along z)
# - apply whitening
# - convert to tensor
val_transforms = Compose([
LoadNiftid(keys=["image"]),
AddChanneld(keys=["image"]),
InPlaneSpacingd(
keys=["image"],
pixdim=spacing,
mode="bilinear",
),
NormalizeIntensityd(keys=["image"], nonzero=False, channel_wise=True),
ToTensord(keys=["image"]),
])
# create a validation data loader
val_ds = Dataset(data=val_files, transform=val_transforms)
val_loader = DataLoader(val_ds,
batch_size=1,
num_workers=config_info['device']['num_workers'])
def prepare_batch(batchdata):
assert isinstance(batchdata,
dict), "prepare_batch expects dictionary input data."
return ((batchdata["image"],
batchdata["label"]) if "label" in batchdata else
(batchdata["image"], None))
"""
Network preparation
"""
print("*** Preparing network ... ")
# automatically extracts the strides and kernels based on nnU-Net empirical rules
spacings = spacing[:2]
sizes = patch_size[:2]
strides, kernels = [], []
while True:
spacing_ratio = [sp / min(spacings) for sp in spacings]
stride = [
2 if ratio <= 2 and size >= 8 else 1
for (ratio, size) in zip(spacing_ratio, sizes)
]
kernel = [3 if ratio <= 2 else 1 for ratio in spacing_ratio]
if all(s == 1 for s in stride):
break
sizes = [i / j for i, j in zip(sizes, stride)]
spacings = [i * j for i, j in zip(spacings, stride)]
kernels.append(kernel)
strides.append(stride)
strides.insert(0, len(spacings) * [1])
kernels.append(len(spacings) * [3])
net = DynUNet(spatial_dims=2,
in_channels=1,
out_channels=nr_out_channels,
kernel_size=kernels,
strides=strides,
upsample_kernel_size=strides[1:],
norm_name="instance",
deep_supervision=True,
deep_supr_num=2,
res_block=False).to(current_device)
"""
Set ignite evaluator to perform inference
"""
print("*** Preparing evaluator ... ")
if nr_out_channels == 1:
do_sigmoid = True
do_softmax = False
elif nr_out_channels > 1:
do_sigmoid = False
do_softmax = True
else:
raise Exception("incompatible number of output channels")
print("Using sigmoid={} and softmax={} as final activation".format(
do_sigmoid, do_softmax))
val_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=do_sigmoid, softmax=do_softmax),
AsDiscreted(keys="pred",
argmax=True,
threshold_values=True,
logit_thresh=prob_thr),
KeepLargestConnectedComponentd(keys="pred", applied_labels=1)
])
val_handlers = [
StatsHandler(output_transform=lambda x: None),
CheckpointLoader(load_path=model_to_load,
load_dict={"net": net},
map_location=torch.device('cpu')),
SegmentationSaver(
output_dir=config_info['output']['out_dir'],
output_ext='.nii.gz',
output_postfix=config_info['output']['out_postfix'],
batch_transform=lambda batch: batch["image_meta_dict"],
output_transform=lambda output: output["pred"],
),
]
# Define customized evaluator
class DynUNetEvaluator(SupervisedEvaluator):
def _iteration(self, engine, batchdata):
inputs, targets = self.prepare_batch(batchdata)
inputs = inputs.to(engine.state.device)
if targets is not None:
targets = targets.to(engine.state.device)
flip_inputs_1 = torch.flip(inputs, dims=(2, ))
flip_inputs_2 = torch.flip(inputs, dims=(3, ))
flip_inputs_3 = torch.flip(inputs, dims=(2, 3))
def _compute_pred():
pred = self.inferer(inputs, self.network)
# use random flipping as data augmentation at inference
flip_pred_1 = torch.flip(self.inferer(flip_inputs_1,
self.network),
dims=(2, ))
flip_pred_2 = torch.flip(self.inferer(flip_inputs_2,
self.network),
dims=(3, ))
flip_pred_3 = torch.flip(self.inferer(flip_inputs_3,
self.network),
dims=(2, 3))
return (pred + flip_pred_1 + flip_pred_2 + flip_pred_3) / 4
# execute forward computation
self.network.eval()
with torch.no_grad():
if self.amp:
with torch.cuda.amp.autocast():
predictions = _compute_pred()
else:
predictions = _compute_pred()
return {"image": inputs, "label": targets, "pred": predictions}
evaluator = DynUNetEvaluator(
device=current_device,
val_data_loader=val_loader,
network=net,
prepare_batch=prepare_batch,
inferer=SlidingWindowInferer2D(roi_size=patch_size,
sw_batch_size=4,
overlap=0.0),
post_transform=val_post_transforms,
val_handlers=val_handlers,
amp=False,
)
"""
Run inference
"""
print("*** Running evaluator ... ")
evaluator.run()
print("Done!")
return
def main_worker(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
if not os.path.exists(args.dir):
raise FileNotFoundError(f"missing directory {args.dir}")
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
# use amp to accelerate training
scaler = torch.cuda.amp.GradScaler()
torch.backends.cudnn.benchmark = True
total_start = time.time()
train_transforms = Compose([
# load 4 Nifti images and stack them together
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(
keys=["image", "label"],
pixdim=(1.0, 1.0, 1.0),
mode=("bilinear", "nearest"),
),
EnsureTyped(keys=["image", "label"]),
ToDeviced(keys=["image", "label"], device=device),
RandSpatialCropd(keys=["image", "label"],
roi_size=[224, 224, 144],
random_size=False),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
])
# create a training data loader
train_ds = BratsCacheDataset(
root_dir=args.dir,
transform=train_transforms,
section="training",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=True,
)
# ThreadDataLoader can be faster if no IO operations when caching all the data in memory
train_loader = ThreadDataLoader(train_ds,
num_workers=0,
batch_size=args.batch_size,
shuffle=True)
# validation transforms and dataset
val_transforms = Compose([
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(
keys=["image", "label"],
pixdim=(1.0, 1.0, 1.0),
mode=("bilinear", "nearest"),
),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
EnsureTyped(keys=["image", "label"]),
ToDeviced(keys=["image", "label"], device=device),
])
val_ds = BratsCacheDataset(
root_dir=args.dir,
transform=val_transforms,
section="validation",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=False,
)
# ThreadDataLoader can be faster if no IO operations when caching all the data in memory
val_loader = ThreadDataLoader(val_ds,
num_workers=0,
batch_size=args.batch_size,
shuffle=False)
# create network, loss function and optimizer
if args.network == "SegResNet":
model = SegResNet(
blocks_down=[1, 2, 2, 4],
blocks_up=[1, 1, 1],
init_filters=16,
in_channels=4,
out_channels=3,
dropout_prob=0.0,
).to(device)
else:
model = UNet(
spatial_dims=3,
in_channels=4,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
loss_function = DiceFocalLoss(
smooth_nr=1e-5,
smooth_dr=1e-5,
squared_pred=True,
to_onehot_y=False,
sigmoid=True,
batch=True,
)
optimizer = Novograd(model.parameters(), lr=args.lr)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=args.epochs)
# wrap the model with DistributedDataParallel module
model = DistributedDataParallel(model, device_ids=[device])
dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True,
reduction="mean_batch")
post_trans = Compose(
[EnsureType(),
Activations(sigmoid=True),
AsDiscrete(threshold=0.5)])
# start a typical PyTorch training
best_metric = -1
best_metric_epoch = -1
print(f"time elapsed before training: {time.time() - total_start}")
train_start = time.time()
for epoch in range(args.epochs):
epoch_start = time.time()
print("-" * 10)
print(f"epoch {epoch + 1}/{args.epochs}")
epoch_loss = train(train_loader, model, loss_function, optimizer,
lr_scheduler, scaler)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
if (epoch + 1) % args.val_interval == 0:
metric, metric_tc, metric_wt, metric_et = evaluate(
model, val_loader, dice_metric, dice_metric_batch, post_trans)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
if dist.get_rank() == 0:
torch.save(model.state_dict(), "best_metric_model.pth")
print(
f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
f" tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f}"
f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
print(
f"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}"
)
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch},"
f" total train time: {(time.time() - train_start):.4f}")
dist.destroy_process_group()
def test_image_normalize_intensityd(self):
key = 'img'
normalizer = NormalizeIntensityd(keys=[key])
normalised = normalizer({key: self.imt})
expected = (self.imt - np.mean(self.imt)) / np.std(self.imt)
self.assertTrue(np.allclose(normalised[key], expected))
