Skip to content

1077 Adjust handler to rank 0 for dynunet pipeline #1078

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 29, 2022
Merged

1077 Adjust handler to rank 0 for dynunet pipeline #1078

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c561e95
adjust handler to rank 0
yiheng-wang-nv Nov 29, 2022
249e6c4
Merge branch 'main' into 1077-modify-handler-in-dynunet
wyli Nov 29, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 26 additions & 42 deletions modules/dynunet_pipeline/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,12 @@
import os
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser

import ignite.distributed as idist
import torch
import torch.distributed as dist
from monai.config import print_config
from monai.handlers import (
CheckpointSaver,
LrScheduleHandler,
MeanDice,
StatsHandler,
ValidationHandler,
from_engine,
)
from monai.handlers import (CheckpointSaver, LrScheduleHandler, MeanDice,
StatsHandler, ValidationHandler, from_engine)
from monai.inferers import SimpleInferer, SlidingWindowInferer
from monai.losses import DiceCELoss
from monai.utils import set_determinism
Expand Down Expand Up @@ -91,6 +86,8 @@ def validation(args):
"mean dice for label {} is {}".format(i + 1, results[:, i].mean())
)

dist.destroy_process_group()


def train(args):
# load hyper parameters
Expand Down Expand Up @@ -151,12 +148,16 @@ def train(args):
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
),
]
val_handlers = (
[
StatsHandler(output_transform=lambda x: None),
CheckpointSaver(
save_dir=val_output_dir, save_dict={"net": net}, save_key_metric=True
),
]
if idist.get_rank() == 0
else None
)

evaluator = DynUNetEvaluator(
device=device,
Expand All @@ -183,16 +184,18 @@ def train(args):

# produce trainer
loss = DiceCELoss(to_onehot_y=True, softmax=True, batch=batch_dice)
train_handlers = []
train_handlers = [
ValidationHandler(validator=evaluator, interval=interval, epoch_level=True)
]
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=from_engine(["loss"], first=True)
),
]
if idist.get_rank() == 0:
train_handlers += [
StatsHandler(
tag_name="train_loss",
output_transform=from_engine(["loss"], first=True),
)
]

trainer = DynUNetTrainer(
device=device,
Expand All @@ -212,27 +215,8 @@ def train(args):
evaluator.logger.setLevel(logging.WARNING)
trainer.logger.setLevel(logging.WARNING)

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)

logger.addHandler(fhandler)

chandler = logging.StreamHandler()
chandler.setLevel(logging.INFO)
chandler.setFormatter(formatter)
logger.addHandler(chandler)

logger.setLevel(logging.INFO)

trainer.run()
dist.destroy_process_group()


if __name__ == "__main__":
Expand Down