def test_one_save_one_load(self):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
net1 = torch.nn.PReLU()
data1 = net1.state_dict()
data1["weight"] = torch.tensor([0.1])
net1.load_state_dict(data1)
net2 = torch.nn.PReLU()
data2 = net2.state_dict()
data2["weight"] = torch.tensor([0.2])
net2.load_state_dict(data2)
engine = Engine(lambda e, b: None)
tempdir = tempfile.mkdtemp()
CheckpointSaver(save_dir=tempdir,
save_dict={
"net": net1
},
save_final=True).attach(engine)
engine.run([0] * 8, max_epochs=5)
path = tempdir + "/net_final_iteration=40.pth"
CheckpointLoader(load_path=path, load_dict={
"net": net2
}).attach(engine)
engine.run([0] * 8, max_epochs=1)
torch.testing.assert_allclose(net2.state_dict()["weight"], 0.1)
shutil.rmtree(tempdir)
def test_partial_over_load(self):
net1 = torch.nn.Sequential(*[torch.nn.PReLU()])
data1 = net1.state_dict()
data1["0.weight"] = torch.tensor([0.1])
net1.load_state_dict(data1)
net2 = torch.nn.Sequential(*[torch.nn.PReLU(), torch.nn.PReLU()])
data2 = net2.state_dict()
data2["0.weight"] = torch.tensor([0.2])
data2["1.weight"] = torch.tensor([0.3])
net2.load_state_dict(data2)
with tempfile.TemporaryDirectory() as tempdir:
engine = Engine(lambda e, b: None)
CheckpointSaver(save_dir=tempdir,
save_dict={
"net": net1
},
save_final=True).attach(engine)
engine.run([0] * 8, max_epochs=5)
path = tempdir + "/net_final_iteration=40.pt"
engine = Engine(lambda e, b: None)
CheckpointLoader(load_path=path,
load_dict={
"net": net2
},
strict=False).attach(engine)
engine.run([0] * 8, max_epochs=1)
torch.testing.assert_allclose(net2.state_dict()["0.weight"].cpu(),
torch.tensor([0.1]))
def test_two_save_one_load(self):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
net1 = torch.nn.PReLU()
optimizer = optim.SGD(net1.parameters(), lr=0.02)
data1 = net1.state_dict()
data1["weight"] = torch.tensor([0.1])
net1.load_state_dict(data1)
net2 = torch.nn.PReLU()
data2 = net2.state_dict()
data2["weight"] = torch.tensor([0.2])
net2.load_state_dict(data2)
with tempfile.TemporaryDirectory() as tempdir:
engine = Engine(lambda e, b: None)
save_dict = {"net": net1, "opt": optimizer}
CheckpointSaver(save_dir=tempdir,
save_dict=save_dict,
save_final=True).attach(engine)
engine.run([0] * 8, max_epochs=5)
path = tempdir + "/checkpoint_final_iteration=40.pt"
engine = Engine(lambda e, b: None)
CheckpointLoader(load_path=path, load_dict={
"net": net2
}).attach(engine)
engine.run([0] * 8, max_epochs=1)
torch.testing.assert_allclose(net2.state_dict()["weight"], 0.1)
def _create_trainer(self, context: Context):
train_handlers: List = self.train_handlers(context)
if context.local_rank == 0:
train_handlers.append(
CheckpointSaver(
save_dir=context.output_dir,
save_dict={self._model_dict_key: context.network},
save_interval=self._train_save_interval,
save_final=True,
final_filename=self._final_filename,
save_key_metric=True,
key_metric_filename=f"train_{self._key_metric_filename}"
if context.evaluator else self._key_metric_filename,
n_saved=5,
))
self._load_checkpoint(context, train_handlers)
return SupervisedTrainer(
device=context.device,
max_epochs=context.max_epochs,
train_data_loader=self.train_data_loader(context),
network=context.network,
optimizer=context.optimizer,
loss_function=self.loss_function(context),
inferer=self.train_inferer(context),
amp=self._amp,
postprocessing=self._validate_transforms(
self.train_post_transforms(context), "Training", "post"),
key_train_metric=self.train_key_metric(context),
train_handlers=train_handlers,
iteration_update=self.train_iteration_update(context),
event_names=self.event_names(context),
)
def test_one_save_one_load(self):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
net1 = torch.nn.PReLU()
data1 = net1.state_dict()
data1["weight"] = torch.tensor([0.1])
net1.load_state_dict(data1)
net2 = torch.nn.PReLU()
data2 = net2.state_dict()
data2["weight"] = torch.tensor([0.2])
net2.load_state_dict(data2)
with tempfile.TemporaryDirectory() as tempdir:
engine1 = Engine(lambda e, b: None)
CheckpointSaver(save_dir=tempdir, save_dict={"net": net1, "eng": engine1}, save_final=True).attach(engine1)
engine1.run([0] * 8, max_epochs=5)
path = tempdir + "/checkpoint_final_iteration=40.pt"
engine2 = Engine(lambda e, b: None)
CheckpointLoader(load_path=path, load_dict={"net": net2, "eng": engine2}).attach(engine2)
@engine2.on(Events.STARTED)
def check_epoch(engine: Engine):
self.assertEqual(engine.state.epoch, 5)
engine2.run([0] * 8, max_epochs=8)
torch.testing.assert_allclose(net2.state_dict()["weight"], torch.tensor([0.1]))
# test bad case with max_epochs smaller than current epoch
engine3 = Engine(lambda e, b: None)
CheckpointLoader(load_path=path, load_dict={"net": net2, "eng": engine3}).attach(engine3)
try:
engine3.run([0] * 8, max_epochs=3)
except ValueError:
self.assertEqual(engine3.state.epoch, 5)
self.assertEqual(engine3.state.max_epochs, 5)
def test_save_single_device_load_multi_devices(self):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
net1 = torch.nn.PReLU()
data1 = net1.state_dict()
data1["weight"] = torch.tensor([0.1])
net1.load_state_dict(data1)
net2 = torch.nn.PReLU()
data2 = net2.state_dict()
data2["weight"] = torch.tensor([0.2])
net2.load_state_dict(data2)
net2 = torch.nn.DataParallel(net2)
with tempfile.TemporaryDirectory() as tempdir:
engine = Engine(lambda e, b: None)
CheckpointSaver(save_dir=tempdir,
save_dict={
"net": net1
},
save_final=True).attach(engine)
engine.run([0] * 8, max_epochs=5)
path = tempdir + "/net_final_iteration=40.pt"
engine = Engine(lambda e, b: None)
CheckpointLoader(load_path=path, load_dict={
"net": net2
}).attach(engine)
engine.run([0] * 8, max_epochs=1)
torch.testing.assert_allclose(
net2.state_dict()["module.weight"].cpu(), 0.1)
def test_strict_shape(self):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
net1 = torch.nn.Sequential(*[torch.nn.PReLU(num_parameters=5)])
data1 = net1.state_dict()
data1["0.weight"] = torch.tensor([1, 2, 3, 4, 5])
data1["new"] = torch.tensor(0.1)
net1.load_state_dict(data1, strict=False)
opt1 = optim.SGD(net1.parameters(), lr=0.02)
net2 = torch.nn.Sequential(*[torch.nn.PReLU(), torch.nn.PReLU()])
data2 = net2.state_dict()
data2["0.weight"] = torch.tensor([0.2])
data2["1.weight"] = torch.tensor([0.3])
net2.load_state_dict(data2)
opt2 = optim.SGD(net2.parameters(), lr=0.02)
with tempfile.TemporaryDirectory() as tempdir:
engine = Engine(lambda e, b: None)
CheckpointSaver(save_dir=tempdir, save_dict={"net": net1, "opt": opt1}, save_final=True).attach(engine)
engine.run([0] * 8, max_epochs=5)
path = tempdir + "/checkpoint_final_iteration=40.pt"
engine = Engine(lambda e, b: None)
CheckpointLoader(
load_path=path,
# expect to print a warning because it loads not only `net` but also `opt` with `strict_shape=False`
load_dict={"net": net2, "opt": opt2},
strict=False,
strict_shape=False,
).attach(engine)
engine.run([0] * 8, max_epochs=1)
torch.testing.assert_allclose(net2.state_dict()["0.weight"].cpu(), torch.tensor([0.2]))
# test whether `opt2` had been skipped when loading with `strict_shape=False`,
# it should have 2 items in `params`(0.weight and 1.weight) while the checkpoint has 1 item(0.weight)
self.assertEqual(len(opt1.state_dict()["param_groups"][0]["params"]), 1)
self.assertEqual(len(opt2.state_dict()["param_groups"][0]["params"]), 2)
def _create_evaluator(self, context: Context):
evaluator = None
if context.val_datalist and len(context.val_datalist) > 0:
val_hanlders: List = self.val_handlers(context)
if context.local_rank == 0:
val_hanlders.append(
CheckpointSaver(
save_dir=context.output_dir,
save_dict={self._model_dict_key: context.network},
save_key_metric=True,
key_metric_filename=self._key_metric_filename,
n_saved=5,
))
evaluator = SupervisedEvaluator(
device=context.device,
val_data_loader=self.val_data_loader(context),
network=context.network,
inferer=self.val_inferer(context),
postprocessing=self._validate_transforms(
self.val_post_transforms(context), "Validation", "post"),
key_val_metric=self.val_key_metric(context),
additional_metrics=self.val_additional_metrics(context),
val_handlers=val_hanlders,
iteration_update=self.val_iteration_update(context),
event_names=self.event_names(context),
)
return evaluator
def test_file(
self,
save_final,
final_filename,
save_key_metric,
key_metric_name,
key_metric_n_saved,
key_metric_filename,
key_metric_save_state,
key_metric_greater_or_equal,
epoch_level,
save_interval,
n_saved,
filenames,
multi_devices=False,
):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
data = [0] * 8
# set up engine
def _train_func(engine, batch):
engine.state.metrics["val_loss"] = engine.state.iteration
engine = Engine(_train_func)
# set up testing handler
net = torch.nn.PReLU()
if multi_devices:
net = torch.nn.DataParallel(net)
optimizer = optim.SGD(net.parameters(), lr=0.02)
with tempfile.TemporaryDirectory() as tempdir:
handler = CheckpointSaver(
tempdir,
{
"net": net,
"opt": optimizer
},
"CheckpointSaver",
"test",
save_final,
final_filename,
save_key_metric,
key_metric_name,
key_metric_n_saved,
key_metric_filename,
key_metric_save_state,
key_metric_greater_or_equal,
epoch_level,
save_interval,
n_saved,
)
handler.attach(engine)
engine.run(data, max_epochs=2)
engine.run(data, max_epochs=5)
for filename in filenames:
self.assertTrue(os.path.exists(os.path.join(tempdir,
filename)))
def test_load_state_dict(self):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
net = torch.nn.PReLU()
# set up engine
def _train_func(engine, batch):
engine.state.metrics["val_loss"] = engine.state.iteration
engine = Engine(_train_func)
# set up testing handler
with tempfile.TemporaryDirectory() as tempdir:
engine = Engine(_train_func)
CheckpointSaver(
save_dir=tempdir,
save_dict={
"net": net
},
save_key_metric=True,
key_metric_name="val_loss",
key_metric_n_saved=2,
key_metric_save_state=True,
).attach(engine)
engine.run(range(3), max_epochs=2)
saver = CheckpointSaver(
save_dir=tempdir,
save_dict={"net": net},
save_key_metric=True,
key_metric_name="val_loss",
key_metric_n_saved=2,
)
engine = Engine(_train_func)
CheckpointLoader(os.path.join(tempdir, "net_key_metric=6.pt"), {
"checkpointer": saver
}).attach(engine)
engine.run(range(1), max_epochs=1)
resumed = saver._key_metric_checkpoint._saved
for i in range(2):
self.assertEqual(resumed[i].priority, 3 * (i + 1))
self.assertEqual(resumed[i].filename,
f"net_key_metric={3 * (i + 1)}.pt")
def run(self, date=None) -> str:
if date is not None:
now = date
else:
now = datetime.datetime.now()
datetime_string = now.strftime('%d/%m/%Y %H:%M:%S')
print(f'Training started: {datetime_string}')
now = datetime.datetime.now()
timedate_info = str(now).split(' ')[0] + '_' + str(now.strftime("%H:%M:%S")).replace(':', '-')
training_dir = os.path.join(self.out_dir, 'training')
if not os.path.exists(training_dir):
os.mkdir(training_dir)
self.output_dir = os.path.join(training_dir, self.out_name + '_' + timedate_info)
os.mkdir(self.output_dir)
self.validator.output_dir = self.output_dir
if self.summary_writer is None:
self.summary_writer = SummaryWriter(log_dir=self.output_dir)
if self.validator.summary_writer is None:
self.validator.summary_writer = self.summary_writer
handlers = [
MetricLogger(self.output_dir, validator=self.validator),
ValidationHandler(
validator=self.validator,
start=self.validation_epoch,
interval=self.validation_interval
),
StatsHandler(tag_name="loss", output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(
summary_writer=self.summary_writer,
tag_name="Loss",
output_transform=lambda x: x["loss"]
),
]
save_dict = { 'network': self.network, 'optimizer': self.optimizer }
if self.lr_scheduler is not None:
handlers.insert(0, LrScheduleHandler(lr_scheduler=self.lr_scheduler, print_lr=True))
save_dict['lr_scheduler'] = self.lr_scheduler
handlers.append(
CheckpointSaver(save_dir=self.output_dir, save_dict=save_dict, save_interval=1, n_saved=1)
)
self._register_handlers(handlers)
super().run()
return self.output_dir
def test_exception(self):
net = torch.nn.PReLU()
# set up engine
def _train_func(engine, batch):
raise RuntimeError("test exception.")
engine = Engine(_train_func)
# set up testing handler
with tempfile.TemporaryDirectory() as tempdir:
stats_handler = CheckpointSaver(tempdir, {"net": net}, save_final=True)
stats_handler.attach(engine)
with self.assertRaises(RuntimeError):
engine.run(range(3), max_epochs=2)
self.assertTrue(os.path.exists(os.path.join(tempdir, "net_final_iteration=1.pt")))
def run(self, global_epoch: int) -> None:
if global_epoch == 1:
handlers = [
StatsHandler(),
TensorBoardStatsHandler(
summary_writer=self.summary_writer
), #, output_transform=lambda x: None),
CheckpointSaver(save_dir=self.output_dir,
save_dict={"network": self.network},
save_key_metric=True),
MetricsSaver(save_dir=self.output_dir,
metrics=['Valid_AUC', 'Valid_ACC']),
self.early_stop_handler,
]
self._register_handlers(handlers)
return super().run(global_epoch=global_epoch)
def main():
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
set_determinism(12345)
device = torch.device("cuda:0")
# load real data
mednist_url = "https://www.dropbox.com/s/5wwskxctvcxiuea/MedNIST.tar.gz?dl=1"
md5_value = "0bc7306e7427e00ad1c5526a6677552d"
extract_dir = "data"
tar_save_path = os.path.join(extract_dir, "MedNIST.tar.gz")
download_and_extract(mednist_url, tar_save_path, extract_dir, md5_value)
hand_dir = os.path.join(extract_dir, "MedNIST", "Hand")
real_data = [{
"hand": os.path.join(hand_dir, filename)
} for filename in os.listdir(hand_dir)]
# define real data transforms
train_transforms = Compose([
LoadPNGD(keys=["hand"]),
AddChannelD(keys=["hand"]),
ScaleIntensityD(keys=["hand"]),
RandRotateD(keys=["hand"], range_x=15, prob=0.5, keep_size=True),
RandFlipD(keys=["hand"], spatial_axis=0, prob=0.5),
RandZoomD(keys=["hand"], min_zoom=0.9, max_zoom=1.1, prob=0.5),
ToTensorD(keys=["hand"]),
])
# create dataset and dataloader
real_dataset = CacheDataset(real_data, train_transforms)
batch_size = 300
real_dataloader = DataLoader(real_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=10)
# define function to process batchdata for input into discriminator
def prepare_batch(batchdata):
"""
Process Dataloader batchdata dict object and return image tensors for D Inferer
"""
return batchdata["hand"]
# define networks
disc_net = Discriminator(in_shape=(1, 64, 64),
channels=(8, 16, 32, 64, 1),
strides=(2, 2, 2, 2, 1),
num_res_units=1,
kernel_size=5).to(device)
latent_size = 64
gen_net = Generator(latent_shape=latent_size,
start_shape=(latent_size, 8, 8),
channels=[32, 16, 8, 1],
strides=[2, 2, 2, 1])
# initialize both networks
disc_net.apply(normal_init)
gen_net.apply(normal_init)
# input images are scaled to [0,1] so enforce the same of generated outputs
gen_net.conv.add_module("activation", torch.nn.Sigmoid())
gen_net = gen_net.to(device)
# create optimizers and loss functions
learning_rate = 2e-4
betas = (0.5, 0.999)
disc_opt = torch.optim.Adam(disc_net.parameters(),
learning_rate,
betas=betas)
gen_opt = torch.optim.Adam(gen_net.parameters(),
learning_rate,
betas=betas)
disc_loss_criterion = torch.nn.BCELoss()
gen_loss_criterion = torch.nn.BCELoss()
real_label = 1
fake_label = 0
def discriminator_loss(gen_images, real_images):
"""
The discriminator loss is calculated by comparing D
prediction for real and generated images.
"""
real = real_images.new_full((real_images.shape[0], 1), real_label)
gen = gen_images.new_full((gen_images.shape[0], 1), fake_label)
realloss = disc_loss_criterion(disc_net(real_images), real)
genloss = disc_loss_criterion(disc_net(gen_images.detach()), gen)
return (genloss + realloss) / 2
def generator_loss(gen_images):
"""
The generator loss is calculated by determining how realistic
the discriminator classifies the generated images.
"""
output = disc_net(gen_images)
cats = output.new_full(output.shape, real_label)
return gen_loss_criterion(output, cats)
# initialize current run dir
run_dir = "model_out"
print("Saving model output to: %s " % run_dir)
# create workflow handlers
handlers = [
StatsHandler(
name="batch_training_loss",
output_transform=lambda x: {
Keys.GLOSS: x[Keys.GLOSS],
Keys.DLOSS: x[Keys.DLOSS]
},
),
CheckpointSaver(
save_dir=run_dir,
save_dict={
"g_net": gen_net,
"d_net": disc_net
},
save_interval=10,
save_final=True,
epoch_level=True,
),
]
# define key metric
key_train_metric = None
# create adversarial trainer
disc_train_steps = 5
num_epochs = 50
trainer = GanTrainer(
device,
num_epochs,
real_dataloader,
gen_net,
gen_opt,
generator_loss,
disc_net,
disc_opt,
discriminator_loss,
d_prepare_batch=prepare_batch,
d_train_steps=disc_train_steps,
latent_shape=latent_size,
key_train_metric=key_train_metric,
train_handlers=handlers,
)
# run GAN training
trainer.run()
# Training completed, save a few random generated images.
print("Saving trained generator sample output.")
test_img_count = 10
test_latents = make_latent(test_img_count, latent_size).to(device)
fakes = gen_net(test_latents)
for i, image in enumerate(fakes):
filename = "gen-fake-final-%d.png" % (i)
save_path = os.path.join(run_dir, filename)
img_array = image[0].cpu().data.numpy()
png_writer.write_png(img_array, save_path, scale=255)
def main(tempdir):
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
################################ DATASET ################################
# create a temporary directory and 40 random image, mask pairs
print(f"generating synthetic data to {tempdir} (this may take a while)")
for i in range(40):
im, seg = create_test_image_3d(128, 128, 128, num_seg_classes=1, channel_dim=-1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(tempdir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(tempdir, f"seg{i:d}.nii.gz"))
images = sorted(glob(os.path.join(tempdir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(tempdir, "seg*.nii.gz")))
train_files = [{"image": img, "label": seg} for img, seg in zip(images[:20], segs[:20])]
val_files = [{"image": img, "label": seg} for img, seg in zip(images[-20:], segs[-20:])]
# define transforms for image and segmentation
train_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys="image"),
RandCropByPosNegLabeld(
keys=["image", "label"], label_key="label", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
),
RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]),
ToTensord(keys=["image", "label"]),
]
)
val_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys="image"),
ToTensord(keys=["image", "label"]),
]
)
# create a training data loader
train_ds = monai.data.CacheDataset(data=train_files, transform=train_transforms, cache_rate=0.5)
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = monai.data.DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)
# create a validation data loader
val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0)
val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
################################ DATASET ################################
################################ NETWORK ################################
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
################################ NETWORK ################################
################################ LOSS ################################
loss = monai.losses.DiceLoss(sigmoid=True)
################################ LOSS ################################
################################ OPT ################################
opt = torch.optim.Adam(net.parameters(), 1e-3)
################################ OPT ################################
################################ LR ################################
lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1)
################################ LR ################################
val_post_transforms = Compose(
[
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
]
)
val_handlers = [
StatsHandler(output_transform=lambda x: None),
TensorBoardStatsHandler(log_dir="./runs/", output_transform=lambda x: None),
TensorBoardImageHandler(
log_dir="./runs/",
batch_transform=lambda x: (x["image"], x["label"]),
output_transform=lambda x: x["pred"],
),
CheckpointSaver(save_dir="./runs/", save_dict={"net": net}, save_key_metric=True),
]
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=SlidingWindowInferer(roi_size=(96, 96, 96), sw_batch_size=4, overlap=0.5),
post_transform=val_post_transforms,
key_val_metric={
"val_mean_dice": MeanDice(include_background=True, output_transform=lambda x: (x["pred"], x["label"]))
},
additional_metrics={"val_acc": Accuracy(output_transform=lambda x: (x["pred"], x["label"]))},
val_handlers=val_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP evaluation
amp=True if monai.utils.get_torch_version_tuple() >= (1, 6) else False,
)
train_post_transforms = Compose(
[
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
]
)
train_handlers = [
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
ValidationHandler(validator=evaluator, interval=2, epoch_level=True),
StatsHandler(tag_name="train_loss", output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(log_dir="./runs/", tag_name="train_loss", output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir="./runs/", save_dict={"net": net, "opt": opt}, save_interval=2, epoch_level=True),
]
trainer = SupervisedTrainer(
device=device,
max_epochs=5,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=loss,
inferer=SimpleInferer(),
post_transform=train_post_transforms,
key_train_metric={"train_acc": Accuracy(output_transform=lambda x: (x["pred"], x["label"]))},
train_handlers=train_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP training
amp=True if monai.utils.get_torch_version_tuple() >= (1, 6) else False,
)
trainer.run()
def train(data_folder=".", model_folder="runs"):
"""run a training pipeline."""
images = sorted(glob.glob(os.path.join(data_folder, "*_ct.nii.gz")))
labels = sorted(glob.glob(os.path.join(data_folder, "*_seg.nii.gz")))
logging.info(
f"training: image/label ({len(images)}) folder: {data_folder}")
amp = True # auto. mixed precision
keys = ("image", "label")
train_frac, val_frac = 0.8, 0.2
n_train = int(train_frac * len(images)) + 1
n_val = min(len(images) - n_train, int(val_frac * len(images)))
logging.info(
f"training: train {n_train} val {n_val}, folder: {data_folder}")
train_files = [{
keys[0]: img,
keys[1]: seg
} for img, seg in zip(images[:n_train], labels[:n_train])]
val_files = [{
keys[0]: img,
keys[1]: seg
} for img, seg in zip(images[-n_val:], labels[-n_val:])]
# create a training data loader
batch_size = 2
logging.info(f"batch size {batch_size}")
train_transforms = get_xforms("train", keys)
train_ds = monai.data.CacheDataset(data=train_files,
transform=train_transforms)
train_loader = monai.data.DataLoader(
train_ds,
batch_size=batch_size,
shuffle=True,
num_workers=2,
pin_memory=torch.cuda.is_available(),
)
# create a validation data loader
val_transforms = get_xforms("val", keys)
val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms)
val_loader = monai.data.DataLoader(
val_ds,
batch_size=
1, # image-level batch to the sliding window method, not the window-level batch
num_workers=2,
pin_memory=torch.cuda.is_available(),
)
# create BasicUNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = get_net().to(device)
max_epochs, lr, momentum = 500, 1e-4, 0.95
logging.info(f"epochs {max_epochs}, lr {lr}, momentum {momentum}")
opt = torch.optim.Adam(net.parameters(), lr=lr)
# create evaluator (to be used to measure model quality during training
val_post_transform = monai.transforms.Compose([
AsDiscreted(keys=("pred", "label"),
argmax=(True, False),
to_onehot=True,
n_classes=2)
])
val_handlers = [
ProgressBar(),
CheckpointSaver(save_dir=model_folder,
save_dict={"net": net},
save_key_metric=True,
key_metric_n_saved=3),
]
evaluator = monai.engines.SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=get_inferer(),
post_transform=val_post_transform,
key_val_metric={
"val_mean_dice":
MeanDice(include_background=False,
output_transform=lambda x: (x["pred"], x["label"]))
},
val_handlers=val_handlers,
amp=amp,
)
# evaluator as an event handler of the trainer
train_handlers = [
ValidationHandler(validator=evaluator, interval=1, epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
]
trainer = monai.engines.SupervisedTrainer(
device=device,
max_epochs=max_epochs,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=DiceCELoss(),
inferer=get_inferer(),
key_train_metric=None,
train_handlers=train_handlers,
amp=amp,
)
trainer.run()
def train(args):
"""run a training pipeline."""
save_args_to_file(args, 'runs/')
images = sorted(glob.glob(os.path.join(args.data_folder, "*_ct.nii.gz")))
labels = sorted(glob.glob(os.path.join(args.data_folder, "*_seg.nii.gz")))
logging.info(
f"training: image/label ({len(images)}) folder: {args.data_folder}")
amp = True # auto. mixed precision
keys = ("image", "label")
#TODO
is_one_hot = False # whether the label has multiple channels to represent multiple class
train_frac, val_frac = 0.8, 0.2
n_train = int(train_frac * len(images)) + 1
n_val = min(len(images) - n_train, int(val_frac * len(images)))
logging.info(
f"training: train {n_train} val {n_val}, folder: {args.data_folder}")
train_files = [{
keys[0]: img,
keys[1]: seg
} for img, seg in zip(images[:n_train], labels[:n_train])]
val_files = [{
keys[0]: img,
keys[1]: seg
} for img, seg in zip(images[-n_val:], labels[-n_val:])]
# create a training data loader
logging.info(f"batch size {args.batch_size}")
train_transforms = get_xforms(args, "train", keys)
train_ds = monai.data.CacheDataset(data=train_files,
transform=train_transforms,
cache_rate=args.cache_rate,
num_workers=args.preprocessing_workers)
train_loader = monai.data.DataLoader(
train_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=torch.cuda.is_available(),
)
# create a validation data loader
val_transforms = get_xforms(args, "val", keys)
val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms)
val_loader = monai.data.DataLoader(
val_ds,
batch_size=
1, # image-level batch to the sliding window method, not the window-level batch
num_workers=args.num_workers,
pin_memory=torch.cuda.is_available(),
)
# create BasicUNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = get_net(args.n_classes).to(device)
logging.info(
f"epochs {args.max_epochs}, lr {args.lr}, momentum {args.momentum}")
opt = torch.optim.Adam(net.parameters(), lr=args.lr)
# create evaluator (to be used to measure model quality during training
def pred_transform(y_pred):
y_sigmoid = torch.sigmoid(y_pred)
y_sigmoid = (y_sigmoid >= logit_thresh).float()
return y_sigmoid
logit_thresh = 0.5
train_metric = MeanDice(
include_background=False,
device=device,
output_transform=lambda x: (pred_transform(x["pred"]), x["label"]),
)
val_metric = MeanDice(
include_background=False,
device=device,
output_transform=lambda x: (pred_transform(x["pred"]), x["label"]),
)
val_handlers = [
ProgressBar(),
CheckpointSaver(save_dir=args.model_folder,
save_dict={
'net': net,
'optimizer': opt
},
save_key_metric=True,
key_metric_n_saved=3),
]
evaluator = monai.engines.SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=get_inferer(args),
key_val_metric={"val_mean_dice": val_metric},
val_handlers=val_handlers,
amp=amp,
)
# evaluator as an event handler of the trainer
train_handlers = [
ValidationHandler(validator=evaluator, interval=1, epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
LrScheduleHandler(
BoundingExponentialLR(opt,
gamma=args.gamma,
min_lr=args.min_lr,
initial_lr=args.lr),
print_lr=True,
name='bounding_lr_scheduler',
epoch_level=True,
)
]
trainer = monai.engines.SupervisedTrainer(
device=device,
max_epochs=args.max_epochs,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=DiceCELoss(),
inferer=get_inferer(args),
key_train_metric={'train_mean_dice': train_metric},
train_handlers=train_handlers,
amp=amp,
)
trainer.run()
def train(gpu, args):
"""run a training pipeline."""
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
print('Setting up multiple GPUs')
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * args.ngpus_per_node + gpu
print(args.rank)
dist.init_process_group(
backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank,
)
print('Done!')
#========================================
images = sorted(glob.glob(os.path.join(args.data_folder, "*_ct.nii.gz")))
labels = sorted(glob.glob(os.path.join(args.data_folder, "*_seg.nii.gz")))
logging.info(f"training: image/label ({len(images)}) folder: {args.data_folder}")
amp = True # auto. mixed precision
keys = ("image", "label")
#TODO
is_one_hot = False # whether the label has multiple channels to represent multiple class
train_frac, val_frac = 0.8, 0.2
n_train = int(train_frac * len(images)) + 1
n_val = min(len(images) - n_train, int(val_frac * len(images)))
logging.info(f"training: train {n_train} val {n_val}, folder: {args.data_folder}")
train_files = [{keys[0]: img, keys[1]: seg} for img, seg in zip(images[:n_train], labels[:n_train])]
val_files = [{keys[0]: img, keys[1]: seg} for img, seg in zip(images[-n_val:], labels[-n_val:])]
# create a training data loader
logging.info(f"batch size {args.batch_size}")
train_transforms = get_xforms("train", keys)
train_ds = monai.data.CacheDataset(data=train_files,
transform=train_transforms,
cache_rate=args.cache_rate,
num_workers=args.preprocessing_workers)
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,
num_replicas=args.world_size,
rank=args.rank
)
train_loader = monai.data.DataLoader(
train_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=torch.cuda.is_available(),
sampler=train_sampler) #
else:
train_loader = monai.data.DataLoader(
train_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=torch.cuda.is_available())
# create a validation data loader
val_transforms = get_xforms("val", keys)
val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms)
val_loader = monai.data.DataLoader(
val_ds,
batch_size=1, # image-level batch to the sliding window method, not the window-level batch
num_workers=args.num_workers,
pin_memory=torch.cuda.is_available(),
)
# create BasicUNet, DiceLoss and Adam optimizer
if args.distributed:
print('Setting Up ')
torch.cuda.set_device(args.gpu)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net.cuda(args.gpu)
args.batch_size = int(args.batch_size / ngpus_per_node)
args.val_batch_size = int(args.val_batch_size / ngpus_per_node)
args.num_workers = int(
(args.num_workers + ngpus_per_node - 1) / ngpus_per_node
)
net = torch.nn.parallel.DistributedDataParallel(
net, device_ids=[args.gpu]
)
else:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = get_net().to(device)
logging.info(f"epochs {args.max_epochs}, lr {args.lr}, momentum {args.momentum}")
opt = torch.optim.Adam(net.parameters(), lr=args.lr)
# create evaluator (to be used to measure model quality during training
def pred_transform(y_pred):
y_sigmoid = torch.sigmoid(y_pred)
y_sigmoid = (y_sigmoid >= logit_thresh).float()
return y_sigmoid
logit_thresh = 0.5
train_metric = MeanDice(
include_background=False,
device = device,
output_transform=lambda x: (pred_transform(x["pred"]), x["label"]),
)
val_metric = MeanDice(
include_background=False,
device = device,
output_transform=lambda x: (pred_transform(x["pred"]), x["label"]),
)
val_handlers = [
ProgressBar(),
CheckpointSaver(save_dir=args.model_folder, save_dict={'net': net,
'optimizer': opt},
save_key_metric=True, key_metric_n_saved=3),
]
evaluator = monai.engines.SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=get_inferer(),
key_val_metric={"val_mean_dice": val_metric},
val_handlers=val_handlers,
amp=amp,
)
# evaluator as an event handler of the trainer
train_handlers = [
ValidationHandler(validator=evaluator, interval=1, epoch_level=True),
StatsHandler(tag_name="train_loss", output_transform=lambda x: x["loss"]),
LrScheduleHandler(BoundingExponentialLR(opt, gamma=args.gamma),
print_lr=True,
name='bounding_lr_scheduler',
epoch_level=True,)
]
trainer = monai.engines.SupervisedTrainer(
device=device,
max_epochs=args.max_epochs,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=DiceCELoss(),
inferer=get_inferer(),
key_train_metric={'train_mean_dice': train_metric},
train_handlers=train_handlers,
amp=amp,
)
trainer.run()
def train(self,
train_info,
valid_info,
hyperparameters,
run_data_check=False):
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
if not run_data_check:
start_dt = datetime.datetime.now()
start_dt_string = start_dt.strftime('%d/%m/%Y %H:%M:%S')
print(f'Training started: {start_dt_string}')
# 1. Create folders to save the model
timedate_info = str(
datetime.datetime.now()).split(' ')[0] + '_' + str(
datetime.datetime.now().strftime("%H:%M:%S")).replace(
':', '-')
path_to_model = os.path.join(
self.out_dir, 'trained_models',
self.unique_name + '_' + timedate_info)
os.mkdir(path_to_model)
# 2. Load hyperparameters
learning_rate = hyperparameters['learning_rate']
weight_decay = hyperparameters['weight_decay']
total_epoch = hyperparameters['total_epoch']
multiplicator = hyperparameters['multiplicator']
batch_size = hyperparameters['batch_size']
validation_epoch = hyperparameters['validation_epoch']
validation_interval = hyperparameters['validation_interval']
H = hyperparameters['H']
L = hyperparameters['L']
# 3. Consider class imbalance
negative, positive = 0, 0
for _, label in train_info:
if int(label) == 0:
negative += 1
elif int(label) == 1:
positive += 1
pos_weight = torch.Tensor([(negative / positive)]).to(self.device)
# 4. Create train and validation loaders, batch_size = 10 for validation loader (10 central slices)
train_data = get_data_from_info(self.image_data_dir, self.seg_data_dir,
train_info)
valid_data = get_data_from_info(self.image_data_dir, self.seg_data_dir,
valid_info)
large_image_splitter(train_data, self.cache_dir)
set_determinism(seed=100)
train_trans, valid_trans = self.transformations(H, L)
train_dataset = PersistentDataset(
data=train_data[:],
transform=train_trans,
cache_dir=self.persistent_dataset_dir)
valid_dataset = PersistentDataset(
data=valid_data[:],
transform=valid_trans,
cache_dir=self.persistent_dataset_dir)
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=self.pin_memory,
num_workers=self.num_workers,
collate_fn=PadListDataCollate(
Method.SYMMETRIC, NumpyPadMode.CONSTANT))
valid_loader = DataLoader(valid_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=self.pin_memory,
num_workers=self.num_workers,
collate_fn=PadListDataCollate(
Method.SYMMETRIC, NumpyPadMode.CONSTANT))
# Perform data checks
if run_data_check:
check_data = monai.utils.misc.first(train_loader)
print(check_data["image"].shape, check_data["label"])
for i in range(batch_size):
multi_slice_viewer(
check_data["image"][i, 0, :, :, :],
check_data["image_meta_dict"]["filename_or_obj"][i])
exit()
"""c = 1
for d in train_loader:
img = d["image"]
seg = d["seg"][0]
seg, _ = nrrd.read(seg)
img_name = d["image_meta_dict"]["filename_or_obj"][0]
print(c, "Name:", img_name, "Size:", img.nelement()*img.element_size()/1024/1024, "MB", "shape:", img.shape)
multi_slice_viewer(img[0, 0, :, :, :], d["image_meta_dict"]["filename_or_obj"][0])
#multi_slice_viewer(seg, d["image_meta_dict"]["filename_or_obj"][0])
c += 1
exit()"""
# 5. Prepare model
model = ModelCT().to(self.device)
# 6. Define loss function, optimizer and scheduler
loss_function = torch.nn.BCEWithLogitsLoss(
pos_weight) # pos_weight for class imbalance
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer,
multiplicator,
last_epoch=-1)
# 7. Create post validation transforms and handlers
path_to_tensorboard = os.path.join(self.out_dir, 'tensorboard')
writer = SummaryWriter(log_dir=path_to_tensorboard)
valid_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
])
valid_handlers = [
StatsHandler(output_transform=lambda x: None),
TensorBoardStatsHandler(summary_writer=writer,
output_transform=lambda x: None),
CheckpointSaver(save_dir=path_to_model,
save_dict={"model": model},
save_key_metric=True),
MetricsSaver(save_dir=path_to_model,
metrics=['Valid_AUC', 'Valid_ACC']),
]
# 8. Create validatior
discrete = AsDiscrete(threshold_values=True)
evaluator = SupervisedEvaluator(
device=self.device,
val_data_loader=valid_loader,
network=model,
post_transform=valid_post_transforms,
key_val_metric={
"Valid_AUC":
ROCAUC(output_transform=lambda x: (x["pred"], x["label"]))
},
additional_metrics={
"Valid_Accuracy":
Accuracy(output_transform=lambda x:
(discrete(x["pred"]), x["label"]))
},
val_handlers=valid_handlers,
amp=self.amp,
)
# 9. Create trainer
# Loss function does the last sigmoid, so we dont need it here.
train_post_transforms = Compose([
# Empty
])
logger = MetricLogger(evaluator=evaluator)
train_handlers = [
logger,
LrScheduleHandler(lr_scheduler=scheduler, print_lr=True),
ValidationHandlerCT(validator=evaluator,
start=validation_epoch,
interval=validation_interval,
epoch_level=True),
StatsHandler(tag_name="loss",
output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(summary_writer=writer,
tag_name="Train_Loss",
output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir=path_to_model,
save_dict={
"model": model,
"opt": optimizer
},
save_interval=1,
n_saved=1),
]
trainer = SupervisedTrainer(
device=self.device,
max_epochs=total_epoch,
train_data_loader=train_loader,
network=model,
optimizer=optimizer,
loss_function=loss_function,
post_transform=train_post_transforms,
train_handlers=train_handlers,
amp=self.amp,
)
# 10. Run trainer
trainer.run()
# 11. Save results
np.save(path_to_model + '/AUCS.npy',
np.array(logger.metrics['Valid_AUC']))
np.save(path_to_model + '/ACCS.npy',
np.array(logger.metrics['Valid_ACC']))
np.save(path_to_model + '/LOSSES.npy', np.array(logger.loss))
np.save(path_to_model + '/PARAMETERS.npy', np.array(hyperparameters))
return path_to_model
def 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
def train(args):
# load hyper parameters
task_id = args.task_id
fold = args.fold
val_output_dir = "./runs_{}_fold{}_{}/".format(task_id, fold,
args.expr_name)
log_filename = "nnunet_task{}_fold{}.log".format(task_id, fold)
log_filename = os.path.join(val_output_dir, log_filename)
interval = args.interval
learning_rate = args.learning_rate
max_epochs = args.max_epochs
multi_gpu_flag = args.multi_gpu
amp_flag = args.amp
lr_decay_flag = args.lr_decay
sw_batch_size = args.sw_batch_size
tta_val = args.tta_val
batch_dice = args.batch_dice
window_mode = args.window_mode
eval_overlap = args.eval_overlap
local_rank = args.local_rank
determinism_flag = args.determinism_flag
determinism_seed = args.determinism_seed
if determinism_flag:
set_determinism(seed=determinism_seed)
if local_rank == 0:
print("Using deterministic training.")
# transforms
train_batch_size = data_loader_params[task_id]["batch_size"]
if multi_gpu_flag:
dist.init_process_group(backend="nccl", init_method="env://")
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
else:
device = torch.device("cuda")
properties, val_loader = get_data(args, mode="validation")
_, train_loader = get_data(args, batch_size=train_batch_size, mode="train")
# produce the network
checkpoint = args.checkpoint
net = get_network(properties, task_id, val_output_dir, checkpoint)
net = net.to(device)
if multi_gpu_flag:
net = DistributedDataParallel(module=net,
device_ids=[device],
find_unused_parameters=True)
optimizer = torch.optim.SGD(
net.parameters(),
lr=learning_rate,
momentum=0.99,
weight_decay=3e-5,
nesterov=True,
)
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer, lr_lambda=lambda epoch: (1 - epoch / max_epochs)**0.9)
# produce evaluator
val_handlers = [
StatsHandler(output_transform=lambda x: None),
CheckpointSaver(save_dir=val_output_dir,
save_dict={"net": net},
save_key_metric=True),
]
evaluator = DynUNetEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
n_classes=len(properties["labels"]),
inferer=SlidingWindowInferer(
roi_size=patch_size[task_id],
sw_batch_size=sw_batch_size,
overlap=eval_overlap,
mode=window_mode,
),
post_transform=None,
key_val_metric={
"val_mean_dice":
MeanDice(
include_background=False,
output_transform=lambda x: (x["pred"], x["label"]),
)
},
val_handlers=val_handlers,
amp=amp_flag,
tta_val=tta_val,
)
# produce trainer
loss = DiceCELoss(to_onehot_y=True, softmax=True, batch=batch_dice)
train_handlers = []
if lr_decay_flag:
train_handlers += [
LrScheduleHandler(lr_scheduler=scheduler, print_lr=True)
]
train_handlers += [
ValidationHandler(validator=evaluator,
interval=interval,
epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
]
trainer = DynUNetTrainer(
device=device,
max_epochs=max_epochs,
train_data_loader=train_loader,
network=net,
optimizer=optimizer,
loss_function=loss,
inferer=SimpleInferer(),
post_transform=None,
key_train_metric=None,
train_handlers=train_handlers,
amp=amp_flag,
)
# run
logger = logging.getLogger()
formatter = logging.Formatter(
"%(asctime)s - %(name)s - %(levelname)s - %(message)s")
# Setup file handler
fhandler = logging.FileHandler(log_filename)
fhandler.setLevel(logging.INFO)
fhandler.setFormatter(formatter)
# Configure stream handler for the cells
chandler = logging.StreamHandler()
chandler.setLevel(logging.INFO)
chandler.setFormatter(formatter)
# Add both handlers
if local_rank == 0:
logger.addHandler(fhandler)
logger.addHandler(chandler)
logger.setLevel(logging.INFO)
trainer.run()
def create_trainer(args):
set_determinism(seed=args.seed)
multi_gpu = args.multi_gpu
local_rank = args.local_rank
if multi_gpu:
dist.init_process_group(backend="nccl", init_method="env://")
device = torch.device("cuda:{}".format(local_rank))
torch.cuda.set_device(device)
else:
device = torch.device("cuda" if args.use_gpu else "cpu")
pre_transforms = get_pre_transforms(args.roi_size, args.model_size,
args.dimensions)
click_transforms = get_click_transforms()
post_transform = get_post_transforms()
train_loader, val_loader = get_loaders(args, pre_transforms)
# define training components
network = get_network(args.network, args.channels,
args.dimensions).to(device)
if multi_gpu:
network = torch.nn.parallel.DistributedDataParallel(
network, device_ids=[local_rank], output_device=local_rank)
if args.resume:
logging.info('{}:: Loading Network...'.format(local_rank))
map_location = {"cuda:0": "cuda:{}".format(local_rank)}
network.load_state_dict(
torch.load(args.model_filepath, map_location=map_location))
# define event-handlers for engine
val_handlers = [
StatsHandler(output_transform=lambda x: None),
TensorBoardStatsHandler(log_dir=args.output,
output_transform=lambda x: None),
DeepgrowStatsHandler(log_dir=args.output,
tag_name='val_dice',
image_interval=args.image_interval),
CheckpointSaver(save_dir=args.output,
save_dict={"net": network},
save_key_metric=True,
save_final=True,
save_interval=args.save_interval,
final_filename='model.pt')
]
val_handlers = val_handlers if local_rank == 0 else None
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=network,
iteration_update=Interaction(
transforms=click_transforms,
max_interactions=args.max_val_interactions,
key_probability='probability',
train=False),
inferer=SimpleInferer(),
post_transform=post_transform,
key_val_metric={
"val_dice":
MeanDice(include_background=False,
output_transform=lambda x: (x["pred"], x["label"]))
},
val_handlers=val_handlers)
loss_function = DiceLoss(sigmoid=True, squared_pred=True)
optimizer = torch.optim.Adam(network.parameters(), args.learning_rate)
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=evaluator,
interval=args.val_freq,
epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(log_dir=args.output,
tag_name="train_loss",
output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir=args.output,
save_dict={
"net": network,
"opt": optimizer,
"lr": lr_scheduler
},
save_interval=args.save_interval * 2,
save_final=True,
final_filename='checkpoint.pt'),
]
train_handlers = train_handlers if local_rank == 0 else train_handlers[:2]
trainer = SupervisedTrainer(
device=device,
max_epochs=args.epochs,
train_data_loader=train_loader,
network=network,
iteration_update=Interaction(
transforms=click_transforms,
max_interactions=args.max_train_interactions,
key_probability='probability',
train=True),
optimizer=optimizer,
loss_function=loss_function,
inferer=SimpleInferer(),
post_transform=post_transform,
amp=args.amp,
key_train_metric={
"train_dice":
MeanDice(include_background=False,
output_transform=lambda x: (x["pred"], x["label"]))
},
train_handlers=train_handlers,
)
return trainer
def run_training_test(root_dir, device="cuda:0", amp=False):
images = sorted(glob(os.path.join(root_dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(root_dir, "seg*.nii.gz")))
train_files = [{
"image": img,
"label": seg
} for img, seg in zip(images[:20], segs[:20])]
val_files = [{
"image": img,
"label": seg
} for img, seg in zip(images[-20:], segs[-20:])]
# define transforms for image and segmentation
train_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys=["image", "label"]),
RandCropByPosNegLabeld(keys=["image", "label"],
label_key="label",
spatial_size=[96, 96, 96],
pos=1,
neg=1,
num_samples=4),
RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]),
ToTensord(keys=["image", "label"]),
])
val_transforms = Compose([
LoadNiftid(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys=["image", "label"]),
ToTensord(keys=["image", "label"]),
])
# create a training data loader
train_ds = monai.data.CacheDataset(data=train_files,
transform=train_transforms,
cache_rate=0.5)
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = monai.data.DataLoader(train_ds,
batch_size=2,
shuffle=True,
num_workers=4)
# create a validation data loader
val_ds = monai.data.CacheDataset(data=val_files,
transform=val_transforms,
cache_rate=1.0)
val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
# create UNet, DiceLoss and Adam optimizer
net = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
loss = monai.losses.DiceLoss(sigmoid=True)
opt = torch.optim.Adam(net.parameters(), 1e-3)
lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1)
val_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
])
val_handlers = [
StatsHandler(output_transform=lambda x: None),
TensorBoardStatsHandler(log_dir=root_dir,
output_transform=lambda x: None),
TensorBoardImageHandler(log_dir=root_dir,
batch_transform=lambda x:
(x["image"], x["label"]),
output_transform=lambda x: x["pred"]),
CheckpointSaver(save_dir=root_dir,
save_dict={"net": net},
save_key_metric=True),
]
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=SlidingWindowInferer(roi_size=(96, 96, 96),
sw_batch_size=4,
overlap=0.5),
post_transform=val_post_transforms,
key_val_metric={
"val_mean_dice":
MeanDice(include_background=True,
output_transform=lambda x: (x["pred"], x["label"]))
},
additional_metrics={
"val_acc":
Accuracy(output_transform=lambda x: (x["pred"], x["label"]))
},
val_handlers=val_handlers,
amp=True if amp else False,
)
train_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
])
train_handlers = [
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
ValidationHandler(validator=evaluator, interval=2, epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(log_dir=root_dir,
tag_name="train_loss",
output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir=root_dir,
save_dict={
"net": net,
"opt": opt
},
save_interval=2,
epoch_level=True),
]
trainer = SupervisedTrainer(
device=device,
max_epochs=5,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=loss,
inferer=SimpleInferer(),
post_transform=train_post_transforms,
key_train_metric={
"train_acc":
Accuracy(output_transform=lambda x: (x["pred"], x["label"]))
},
train_handlers=train_handlers,
amp=True if amp else False,
)
trainer.run()
return evaluator.state.best_metric
def train(cfg):
log_dir = create_log_dir(cfg)
device = set_device(cfg)
# --------------------------------------------------------------------------
# Data Loading and Preprocessing
# --------------------------------------------------------------------------
# __________________________________________________________________________
# Build MONAI preprocessing
train_preprocess = Compose([
ToTensorD(keys="image"),
TorchVisionD(keys="image",
name="ColorJitter",
brightness=64.0 / 255.0,
contrast=0.75,
saturation=0.25,
hue=0.04),
ToNumpyD(keys="image"),
RandFlipD(keys="image", prob=0.5),
RandRotate90D(keys="image", prob=0.5),
CastToTypeD(keys="image", dtype=np.float32),
RandZoomD(keys="image", prob=0.5, min_zoom=0.9, max_zoom=1.1),
ScaleIntensityRangeD(keys="image",
a_min=0.0,
a_max=255.0,
b_min=-1.0,
b_max=1.0),
ToTensorD(keys=("image", "label")),
])
valid_preprocess = Compose([
CastToTypeD(keys="image", dtype=np.float32),
ScaleIntensityRangeD(keys="image",
a_min=0.0,
a_max=255.0,
b_min=-1.0,
b_max=1.0),
ToTensorD(keys=("image", "label")),
])
# __________________________________________________________________________
# Create MONAI dataset
train_json_info_list = load_decathlon_datalist(
data_list_file_path=cfg["dataset_json"],
data_list_key="training",
base_dir=cfg["data_root"],
)
valid_json_info_list = load_decathlon_datalist(
data_list_file_path=cfg["dataset_json"],
data_list_key="validation",
base_dir=cfg["data_root"],
)
train_dataset = PatchWSIDataset(
train_json_info_list,
cfg["region_size"],
cfg["grid_shape"],
cfg["patch_size"],
train_preprocess,
image_reader_name="openslide" if cfg["use_openslide"] else "cuCIM",
)
valid_dataset = PatchWSIDataset(
valid_json_info_list,
cfg["region_size"],
cfg["grid_shape"],
cfg["patch_size"],
valid_preprocess,
image_reader_name="openslide" if cfg["use_openslide"] else "cuCIM",
)
# __________________________________________________________________________
# DataLoaders
train_dataloader = DataLoader(train_dataset,
num_workers=cfg["num_workers"],
batch_size=cfg["batch_size"],
pin_memory=True)
valid_dataloader = DataLoader(valid_dataset,
num_workers=cfg["num_workers"],
batch_size=cfg["batch_size"],
pin_memory=True)
# __________________________________________________________________________
# Get sample batch and some info
first_sample = first(train_dataloader)
if first_sample is None:
raise ValueError("Fist sample is None!")
print("image: ")
print(" shape", first_sample["image"].shape)
print(" type: ", type(first_sample["image"]))
print(" dtype: ", first_sample["image"].dtype)
print("labels: ")
print(" shape", first_sample["label"].shape)
print(" type: ", type(first_sample["label"]))
print(" dtype: ", first_sample["label"].dtype)
print(f"batch size: {cfg['batch_size']}")
print(f"train number of batches: {len(train_dataloader)}")
print(f"valid number of batches: {len(valid_dataloader)}")
# --------------------------------------------------------------------------
# Deep Learning Classification Model
# --------------------------------------------------------------------------
# __________________________________________________________________________
# initialize model
model = TorchVisionFCModel("resnet18",
num_classes=1,
use_conv=True,
pretrained=cfg["pretrain"])
model = model.to(device)
# loss function
loss_func = torch.nn.BCEWithLogitsLoss()
loss_func = loss_func.to(device)
# optimizer
if cfg["novograd"]:
optimizer = Novograd(model.parameters(), cfg["lr"])
else:
optimizer = SGD(model.parameters(), lr=cfg["lr"], momentum=0.9)
# AMP scaler
if cfg["amp"]:
cfg["amp"] = True if monai.utils.get_torch_version_tuple() >= (
1, 6) else False
else:
cfg["amp"] = False
scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
T_max=cfg["n_epochs"])
# --------------------------------------------
# Ignite Trainer/Evaluator
# --------------------------------------------
# Evaluator
val_handlers = [
CheckpointSaver(save_dir=log_dir,
save_dict={"net": model},
save_key_metric=True),
StatsHandler(output_transform=lambda x: None),
TensorBoardStatsHandler(log_dir=log_dir,
output_transform=lambda x: None),
]
val_postprocessing = Compose([
ActivationsD(keys="pred", sigmoid=True),
AsDiscreteD(keys="pred", threshold=0.5)
])
evaluator = SupervisedEvaluator(
device=device,
val_data_loader=valid_dataloader,
network=model,
postprocessing=val_postprocessing,
key_val_metric={
"val_acc":
Accuracy(output_transform=from_engine(["pred", "label"]))
},
val_handlers=val_handlers,
amp=cfg["amp"],
)
# Trainer
train_handlers = [
LrScheduleHandler(lr_scheduler=scheduler, print_lr=True),
CheckpointSaver(save_dir=cfg["logdir"],
save_dict={
"net": model,
"opt": optimizer
},
save_interval=1,
epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=from_engine(["loss"], first=True)),
ValidationHandler(validator=evaluator, interval=1, epoch_level=True),
TensorBoardStatsHandler(log_dir=cfg["logdir"],
tag_name="train_loss",
output_transform=from_engine(["loss"],
first=True)),
]
train_postprocessing = Compose([
ActivationsD(keys="pred", sigmoid=True),
AsDiscreteD(keys="pred", threshold=0.5)
])
trainer = SupervisedTrainer(
device=device,
max_epochs=cfg["n_epochs"],
train_data_loader=train_dataloader,
network=model,
optimizer=optimizer,
loss_function=loss_func,
postprocessing=train_postprocessing,
key_train_metric={
"train_acc":
Accuracy(output_transform=from_engine(["pred", "label"]))
},
train_handlers=train_handlers,
amp=cfg["amp"],
)
trainer.run()
def run_training_test(root_dir, device="cuda:0"):
real_images = sorted(glob(os.path.join(root_dir, "img*.nii.gz")))
train_files = [{"reals": img} for img in zip(real_images)]
# prepare real data
train_transforms = Compose([
LoadNiftid(keys=["reals"]),
AsChannelFirstd(keys=["reals"]),
ScaleIntensityd(keys=["reals"]),
RandFlipd(keys=["reals"], prob=0.5),
ToTensord(keys=["reals"]),
])
train_ds = monai.data.CacheDataset(data=train_files,
transform=train_transforms,
cache_rate=0.5)
train_loader = monai.data.DataLoader(train_ds,
batch_size=2,
shuffle=True,
num_workers=4)
learning_rate = 2e-4
betas = (0.5, 0.999)
real_label = 1
fake_label = 0
# create discriminator
disc_net = Discriminator(in_shape=(1, 64, 64),
channels=(8, 16, 32, 64, 1),
strides=(2, 2, 2, 2, 1),
num_res_units=1,
kernel_size=5).to(device)
disc_net.apply(normal_init)
disc_opt = torch.optim.Adam(disc_net.parameters(),
learning_rate,
betas=betas)
disc_loss_criterion = torch.nn.BCELoss()
def discriminator_loss(gen_images, real_images):
real = real_images.new_full((real_images.shape[0], 1), real_label)
gen = gen_images.new_full((gen_images.shape[0], 1), fake_label)
realloss = disc_loss_criterion(disc_net(real_images), real)
genloss = disc_loss_criterion(disc_net(gen_images.detach()), gen)
return torch.div(torch.add(realloss, genloss), 2)
# create generator
latent_size = 64
gen_net = Generator(latent_shape=latent_size,
start_shape=(latent_size, 8, 8),
channels=[32, 16, 8, 1],
strides=[2, 2, 2, 1])
gen_net.apply(normal_init)
gen_net.conv.add_module("activation", torch.nn.Sigmoid())
gen_net = gen_net.to(device)
gen_opt = torch.optim.Adam(gen_net.parameters(),
learning_rate,
betas=betas)
gen_loss_criterion = torch.nn.BCELoss()
def generator_loss(gen_images):
output = disc_net(gen_images)
cats = output.new_full(output.shape, real_label)
return gen_loss_criterion(output, cats)
key_train_metric = None
train_handlers = [
StatsHandler(
name="training_loss",
output_transform=lambda x: {
Keys.GLOSS: x[Keys.GLOSS],
Keys.DLOSS: x[Keys.DLOSS]
},
),
TensorBoardStatsHandler(
log_dir=root_dir,
tag_name="training_loss",
output_transform=lambda x: {
Keys.GLOSS: x[Keys.GLOSS],
Keys.DLOSS: x[Keys.DLOSS]
},
),
CheckpointSaver(save_dir=root_dir,
save_dict={
"g_net": gen_net,
"d_net": disc_net
},
save_interval=2,
epoch_level=True),
]
disc_train_steps = 2
num_epochs = 5
trainer = GanTrainer(
device,
num_epochs,
train_loader,
gen_net,
gen_opt,
generator_loss,
disc_net,
disc_opt,
discriminator_loss,
d_train_steps=disc_train_steps,
latent_shape=latent_size,
key_train_metric=key_train_metric,
train_handlers=train_handlers,
)
trainer.run()
return trainer.state
def train(args):
if args.local_rank == 0 and not os.path.exists(args.dir):
# create 40 random image, mask paris for training
print(
f"generating synthetic data to {args.dir} (this may take a while)")
os.makedirs(args.dir)
# set random seed to generate same random data for every node
np.random.seed(seed=0)
for i in range(40):
im, seg = create_test_image_3d(128,
128,
128,
num_seg_classes=1,
channel_dim=-1)
n = nib.Nifti1Image(im, np.eye(4))
nib.save(n, os.path.join(args.dir, f"img{i:d}.nii.gz"))
n = nib.Nifti1Image(seg, np.eye(4))
nib.save(n, os.path.join(args.dir, f"seg{i:d}.nii.gz"))
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
images = sorted(glob(os.path.join(args.dir, "img*.nii.gz")))
segs = sorted(glob(os.path.join(args.dir, "seg*.nii.gz")))
train_files = [{
"image": img,
"label": seg
} for img, seg in zip(images, segs)]
# define transforms for image and segmentation
train_transforms = Compose([
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
ScaleIntensityd(keys="image"),
RandCropByPosNegLabeld(keys=["image", "label"],
label_key="label",
spatial_size=[96, 96, 96],
pos=1,
neg=1,
num_samples=4),
RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]),
ToTensord(keys=["image", "label"]),
])
# create a training data loader
train_ds = Dataset(data=train_files, transform=train_transforms)
# create a training data sampler
train_sampler = DistributedSampler(train_ds)
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = DataLoader(
train_ds,
batch_size=2,
shuffle=False,
num_workers=2,
pin_memory=True,
sampler=train_sampler,
)
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
net = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
loss = monai.losses.DiceLoss(sigmoid=True)
opt = torch.optim.Adam(net.parameters(), 1e-3)
lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1)
# wrap the model with DistributedDataParallel module
net = DistributedDataParallel(net, device_ids=[device])
train_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
])
train_handlers = [
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
]
if dist.get_rank() == 0:
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
train_handlers.extend([
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir="./runs/",
save_dict={
"net": net,
"opt": opt
},
save_interval=2),
])
trainer = SupervisedTrainer(
device=device,
max_epochs=5,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=loss,
inferer=SimpleInferer(),
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP evaluation
amp=True if monai.config.get_torch_version_tuple() >=
(1, 6) else False,
post_transform=train_post_transforms,
key_train_metric={
"train_acc":
Accuracy(output_transform=lambda x: (x["pred"], x["label"]),
device=device)
},
train_handlers=train_handlers,
)
trainer.run()
dist.destroy_process_group()
def configure(self):
self.set_device()
network = UNet(
dimensions=3,
in_channels=1,
out_channels=2,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
).to(self.device)
if self.multi_gpu:
network = DistributedDataParallel(
module=network,
device_ids=[self.device],
find_unused_parameters=False,
)
train_transforms = Compose([
LoadImaged(keys=("image", "label")),
EnsureChannelFirstd(keys=("image", "label")),
Spacingd(keys=("image", "label"),
pixdim=[1.0, 1.0, 1.0],
mode=["bilinear", "nearest"]),
ScaleIntensityRanged(
keys="image",
a_min=-57,
a_max=164,
b_min=0.0,
b_max=1.0,
clip=True,
),
CropForegroundd(keys=("image", "label"), source_key="image"),
RandCropByPosNegLabeld(
keys=("image", "label"),
label_key="label",
spatial_size=(96, 96, 96),
pos=1,
neg=1,
num_samples=4,
image_key="image",
image_threshold=0,
),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
ToTensord(keys=("image", "label")),
])
train_datalist = load_decathlon_datalist(self.data_list_file_path,
True, "training")
if self.multi_gpu:
train_datalist = partition_dataset(
data=train_datalist,
shuffle=True,
num_partitions=dist.get_world_size(),
even_divisible=True,
)[dist.get_rank()]
train_ds = CacheDataset(
data=train_datalist,
transform=train_transforms,
cache_num=32,
cache_rate=1.0,
num_workers=4,
)
train_data_loader = DataLoader(
train_ds,
batch_size=2,
shuffle=True,
num_workers=4,
)
val_transforms = Compose([
LoadImaged(keys=("image", "label")),
EnsureChannelFirstd(keys=("image", "label")),
ScaleIntensityRanged(
keys="image",
a_min=-57,
a_max=164,
b_min=0.0,
b_max=1.0,
clip=True,
),
CropForegroundd(keys=("image", "label"), source_key="image"),
ToTensord(keys=("image", "label")),
])
val_datalist = load_decathlon_datalist(self.data_list_file_path, True,
"validation")
val_ds = CacheDataset(val_datalist, val_transforms, 9, 0.0, 4)
val_data_loader = DataLoader(
val_ds,
batch_size=1,
shuffle=False,
num_workers=4,
)
post_transform = Compose([
Activationsd(keys="pred", softmax=True),
AsDiscreted(
keys=["pred", "label"],
argmax=[True, False],
to_onehot=True,
n_classes=2,
),
])
# metric
key_val_metric = {
"val_mean_dice":
MeanDice(
include_background=False,
output_transform=lambda x: (x["pred"], x["label"]),
device=self.device,
)
}
val_handlers = [
StatsHandler(output_transform=lambda x: None),
CheckpointSaver(
save_dir=self.ckpt_dir,
save_dict={"model": network},
save_key_metric=True,
),
TensorBoardStatsHandler(log_dir=self.ckpt_dir,
output_transform=lambda x: None),
]
self.eval_engine = SupervisedEvaluator(
device=self.device,
val_data_loader=val_data_loader,
network=network,
inferer=SlidingWindowInferer(
roi_size=[160, 160, 160],
sw_batch_size=4,
overlap=0.5,
),
post_transform=post_transform,
key_val_metric=key_val_metric,
val_handlers=val_handlers,
amp=self.amp,
)
optimizer = torch.optim.Adam(network.parameters(), self.learning_rate)
loss_function = DiceLoss(to_onehot_y=True, softmax=True)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=5000,
gamma=0.1)
train_handlers = [
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
ValidationHandler(validator=self.eval_engine,
interval=self.val_interval,
epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(
log_dir=self.ckpt_dir,
tag_name="train_loss",
output_transform=lambda x: x["loss"],
),
]
self.train_engine = SupervisedTrainer(
device=self.device,
max_epochs=self.max_epochs,
train_data_loader=train_data_loader,
network=network,
optimizer=optimizer,
loss_function=loss_function,
inferer=SimpleInferer(),
post_transform=post_transform,
key_train_metric=None,
train_handlers=train_handlers,
amp=self.amp,
)
if self.local_rank > 0:
self.train_engine.logger.setLevel(logging.WARNING)
self.eval_engine.logger.setLevel(logging.WARNING)
def train(data_folder=".", model_folder="runs", continue_training=False):
"""run a training pipeline."""
#/== files for synthesis
path_parent = Path(
'/content/drive/My Drive/Datasets/covid19/COVID-19-20_augs_cea/')
path_synthesis = Path(
path_parent /
'CeA_BASE_grow=1_bg=-1.00_step=-1.0_scale=-1.0_seed=1.0_ch0_1=-1_ch1_16=-1_ali_thr=0.1'
)
scans_syns = os.listdir(path_synthesis)
decreasing_sequence = get_decreasing_sequence(255, splits=20)
keys2 = ("image", "label", "synthetic_lesion")
# READ THE SYTHETIC HEALTHY TEXTURE
path_synthesis_old = '/content/drive/My Drive/Datasets/covid19/results/cea_synthesis/patient0/'
texture_orig = np.load(f'{path_synthesis_old}texture.npy.npz')
texture_orig = texture_orig.f.arr_0
texture = texture_orig + np.abs(np.min(texture_orig)) + .07
texture = np.pad(texture, ((100, 100), (100, 100)), mode='reflect')
print(f'type(texture) = {type(texture)}, {np.shape(texture)}')
#==/
images = sorted(glob.glob(os.path.join(data_folder,
"*_ct.nii.gz"))[:10]) #OMM
labels = sorted(glob.glob(os.path.join(data_folder,
"*_seg.nii.gz"))[:10]) #OMM
logging.info(
f"training: image/label ({len(images)}) folder: {data_folder}")
amp = True # auto. mixed precision
keys = ("image", "label")
train_frac, val_frac = 0.8, 0.2
n_train = int(train_frac * len(images)) + 1
n_val = min(len(images) - n_train, int(val_frac * len(images)))
logging.info(
f"training: train {n_train} val {n_val}, folder: {data_folder}")
train_files = [{
keys[0]: img,
keys[1]: seg
} for img, seg in zip(images[:n_train], labels[:n_train])]
val_files = [{
keys[0]: img,
keys[1]: seg
} for img, seg in zip(images[-n_val:], labels[-n_val:])]
# create a training data loader
batch_size = 1 # XX was 2
logging.info(f"batch size {batch_size}")
train_transforms = get_xforms("synthesis", keys, keys2, path_synthesis,
decreasing_sequence, scans_syns, texture)
train_ds = monai.data.CacheDataset(data=train_files,
transform=train_transforms)
train_loader = monai.data.DataLoader(
train_ds,
batch_size=batch_size,
shuffle=True,
num_workers=2,
pin_memory=torch.cuda.is_available(),
# collate_fn=pad_list_data_collate,
)
# create a validation data loader
val_transforms = get_xforms("val", keys)
val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms)
val_loader = monai.data.DataLoader(
val_ds,
batch_size=
1, # image-level batch to the sliding window method, not the window-level batch
num_workers=2,
pin_memory=torch.cuda.is_available(),
)
# create BasicUNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = get_net().to(device)
# if continue training
if continue_training:
ckpts = sorted(glob.glob(os.path.join(model_folder, "*.pt")))
ckpt = ckpts[-1]
logging.info(f"continue training using {ckpt}.")
net.load_state_dict(torch.load(ckpt, map_location=device))
# max_epochs, lr, momentum = 500, 1e-4, 0.95
max_epochs, lr, momentum = 20, 1e-4, 0.95 #OMM
logging.info(f"epochs {max_epochs}, lr {lr}, momentum {momentum}")
opt = torch.optim.Adam(net.parameters(), lr=lr)
# create evaluator (to be used to measure model quality during training
val_post_transform = monai.transforms.Compose([
AsDiscreted(keys=("pred", "label"),
argmax=(True, False),
to_onehot=True,
n_classes=2)
])
val_handlers = [
ProgressBar(),
MetricsSaver(save_dir="./metrics_val", metrics="*"),
CheckpointSaver(save_dir=model_folder,
save_dict={"net": net},
save_key_metric=True,
key_metric_n_saved=6),
]
evaluator = monai.engines.SupervisedEvaluator(
device=device,
val_data_loader=val_loader,
network=net,
inferer=get_inferer(),
post_transform=val_post_transform,
key_val_metric={
"val_mean_dice":
MeanDice(include_background=False,
output_transform=lambda x: (x["pred"], x["label"]))
},
val_handlers=val_handlers,
amp=amp,
)
# evaluator as an event handler of the trainer
train_handlers = [
ValidationHandler(validator=evaluator, interval=1, epoch_level=True),
# MetricsSaver(save_dir="./metrics_train", metrics="*"),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
]
trainer = monai.engines.SupervisedTrainer(
device=device,
max_epochs=max_epochs,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=DiceCELoss(),
inferer=get_inferer(),
key_train_metric=None,
train_handlers=train_handlers,
amp=amp,
)
trainer.run()
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(tempdir):
monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
################################ DATASET ################################
# get dataset
train_ds = CacheDataset(data=train_files,
transform=train_transforms,
cache_rate=0.5)
train_loader = DataLoader(train_ds,
batch_size=2,
shuffle=True,
num_workers=4)
val_ds = CacheDataset(data=val_files,
transform=val_transforms,
cache_rate=1.0)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)
################################ DATASET ################################
################################ NETWORK ################################
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = monai.networks.nets.UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
################################ NETWORK ################################
################################ LOSS ################################
loss = monai.losses.DiceLoss(sigmoid=True)
################################ LOSS ################################
################################ OPT ################################
opt = torch.optim.Adam(net.parameters(), 1e-3)
################################ OPT ################################
################################ LR ################################
lr_scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1)
################################ LR ################################
################################ Evalutaion ################################
val_post_transforms = ...
val_handlers = ...
evaluator = ...
train_post_transforms = Compose([
Activationsd(keys="pred", sigmoid=True),
AsDiscreted(keys="pred", threshold_values=True),
KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
])
train_handlers = [
LrScheduleHandler(lr_scheduler=lr_scheduler, print_lr=True),
ValidationHandler(validator=evaluator, interval=2, epoch_level=True),
StatsHandler(tag_name="train_loss",
output_transform=lambda x: x["loss"]),
TensorBoardStatsHandler(log_dir="./runs/",
tag_name="train_loss",
output_transform=lambda x: x["loss"]),
CheckpointSaver(save_dir="./runs/",
save_dict={
"net": net,
"opt": opt
},
save_interval=2,
epoch_level=True),
]
trainer = SupervisedTrainer(
device=device,
max_epochs=5,
train_data_loader=train_loader,
network=net,
optimizer=opt,
loss_function=loss,
inferer=SimpleInferer(),
post_transform=train_post_transforms,
key_train_metric={
"train_acc":
Accuracy(output_transform=lambda x: (x["pred"], x["label"]))
},
train_handlers=train_handlers,
# if no FP16 support in GPU or PyTorch version < 1.6, will not enable AMP training
amp=True if monai.utils.get_torch_version_tuple() >= (1, 6) else False,
)
trainer.run()
