def get_task_transforms(mode, task_id, pos_sample_num, neg_sample_num,
num_samples):
if mode != "test":
keys = ["image", "label"]
else:
keys = ["image"]
load_transforms = [
LoadImaged(keys=keys),
EnsureChannelFirstd(keys=keys),
]
# 2. sampling
sample_transforms = [
PreprocessAnisotropic(
keys=keys,
clip_values=clip_values[task_id],
pixdim=spacing[task_id],
normalize_values=normalize_values[task_id],
model_mode=mode,
),
]
# 3. spatial transforms
if mode == "train":
other_transforms = [
SpatialPadd(keys=["image", "label"],
spatial_size=patch_size[task_id]),
RandCropByPosNegLabeld(
keys=["image", "label"],
label_key="label",
spatial_size=patch_size[task_id],
pos=pos_sample_num,
neg=neg_sample_num,
num_samples=num_samples,
image_key="image",
image_threshold=0,
),
RandZoomd(
keys=["image", "label"],
min_zoom=0.9,
max_zoom=1.2,
mode=("trilinear", "nearest"),
align_corners=(True, None),
prob=0.15,
),
RandGaussianNoised(keys=["image"], std=0.01, prob=0.15),
RandGaussianSmoothd(
keys=["image"],
sigma_x=(0.5, 1.15),
sigma_y=(0.5, 1.15),
sigma_z=(0.5, 1.15),
prob=0.15,
),
RandScaleIntensityd(keys=["image"], factors=0.3, prob=0.15),
RandFlipd(["image", "label"], spatial_axis=[0], prob=0.5),
RandFlipd(["image", "label"], spatial_axis=[1], prob=0.5),
RandFlipd(["image", "label"], spatial_axis=[2], prob=0.5),
CastToTyped(keys=["image", "label"], dtype=(np.float32, np.uint8)),
EnsureTyped(keys=["image", "label"]),
]
elif mode == "validation":
other_transforms = [
CastToTyped(keys=["image", "label"], dtype=(np.float32, np.uint8)),
EnsureTyped(keys=["image", "label"]),
]
else:
other_transforms = [
CastToTyped(keys=["image"], dtype=(np.float32)),
EnsureTyped(keys=["image"]),
]
all_transforms = load_transforms + sample_transforms + other_transforms
return Compose(all_transforms)
def main_worker(args):
# disable logging for processes except 0 on every node
if args.local_rank != 0:
f = open(os.devnull, "w")
sys.stdout = sys.stderr = f
if not os.path.exists(args.dir):
raise FileNotFoundError(f"Missing directory {args.dir}")
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
total_start = time.time()
train_transforms = Compose([
# load 4 Nifti images and stack them together
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Spacingd(keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest")),
Orientationd(keys=["image", "label"], axcodes="RAS"),
RandSpatialCropd(keys=["image", "label"],
roi_size=[128, 128, 64],
random_size=False),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
ToTensord(keys=["image", "label"]),
])
# create a training data loader
train_ds = BratsCacheDataset(
root_dir=args.dir,
transform=train_transforms,
section="training",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=True,
)
train_loader = DataLoader(train_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
# validation transforms and dataset
val_transforms = Compose([
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Spacingd(keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest")),
Orientationd(keys=["image", "label"], axcodes="RAS"),
CenterSpatialCropd(keys=["image", "label"], roi_size=[128, 128, 64]),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
ToTensord(keys=["image", "label"]),
])
val_ds = BratsCacheDataset(
root_dir=args.dir,
transform=val_transforms,
section="validation",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=False,
)
val_loader = DataLoader(val_ds,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if dist.get_rank() == 0:
# Logging for TensorBoard
writer = SummaryWriter(log_dir=args.log_dir)
# create UNet, DiceLoss and Adam optimizer
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
if args.network == "UNet":
model = UNet(
dimensions=3,
in_channels=4,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
else:
model = SegResNet(in_channels=4,
out_channels=3,
init_filters=16,
dropout_prob=0.2).to(device)
loss_function = DiceLoss(to_onehot_y=False,
sigmoid=True,
squared_pred=True)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=1e-5,
amsgrad=True)
# wrap the model with DistributedDataParallel module
model = DistributedDataParallel(model, device_ids=[device])
# start a typical PyTorch training
total_epoch = args.epochs
best_metric = -1000000
best_metric_epoch = -1
epoch_time = AverageMeter("Time", ":6.3f")
progress = ProgressMeter(total_epoch, [epoch_time], prefix="Epoch: ")
end = time.time()
print(f"Time elapsed before training: {end-total_start}")
for epoch in range(total_epoch):
train_loss = train(train_loader, model, loss_function, optimizer,
epoch, args, device)
epoch_time.update(time.time() - end)
if epoch % args.print_freq == 0:
progress.display(epoch)
if dist.get_rank() == 0:
writer.add_scalar("Loss/train", train_loss, epoch)
if (epoch + 1) % args.val_interval == 0:
metric, metric_tc, metric_wt, metric_et = evaluate(
model, val_loader, device)
if dist.get_rank() == 0:
writer.add_scalar("Mean Dice/val", metric, epoch)
writer.add_scalar("Mean Dice TC/val", metric_tc, epoch)
writer.add_scalar("Mean Dice WT/val", metric_wt, epoch)
writer.add_scalar("Mean Dice ET/val", metric_et, epoch)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
print(
f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
f" tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f}"
f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
end = time.time()
print(f"Time elapsed after epoch {epoch + 1} is {end - total_start}")
if dist.get_rank() == 0:
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
# all processes should see same parameters as they all start from same
# random parameters and gradients are synchronized in backward passes,
# therefore, saving it in one process is sufficient
torch.save(model.state_dict(), "final_model.pth")
writer.flush()
dist.destroy_process_group()
def main_worker(args):
# disable logging for processes except 0 on every node
if args.local_rank != 0:
f = open(os.devnull, "w")
sys.stdout = sys.stderr = f
if not os.path.exists(args.dir):
raise FileNotFoundError(f"missing directory {args.dir}")
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
# use amp to accelerate training
scaler = torch.cuda.amp.GradScaler()
torch.backends.cudnn.benchmark = True
total_start = time.time()
train_transforms = Compose([
# load 4 Nifti images and stack them together
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(
keys=["image", "label"],
pixdim=(1.0, 1.0, 1.0),
mode=("bilinear", "nearest"),
),
EnsureTyped(keys=["image", "label"]),
ToDeviced(keys=["image", "label"], device=device),
RandSpatialCropd(keys=["image", "label"],
roi_size=[224, 224, 144],
random_size=False),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
])
# create a training data loader
train_ds = BratsCacheDataset(
root_dir=args.dir,
transform=train_transforms,
section="training",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=True,
)
# ThreadDataLoader can be faster if no IO operations when caching all the data in memory
train_loader = ThreadDataLoader(train_ds,
num_workers=0,
batch_size=args.batch_size,
shuffle=True)
# validation transforms and dataset
val_transforms = Compose([
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(
keys=["image", "label"],
pixdim=(1.0, 1.0, 1.0),
mode=("bilinear", "nearest"),
),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
EnsureTyped(keys=["image", "label"]),
ToDeviced(keys=["image", "label"], device=device),
])
val_ds = BratsCacheDataset(
root_dir=args.dir,
transform=val_transforms,
section="validation",
num_workers=4,
cache_rate=args.cache_rate,
shuffle=False,
)
# ThreadDataLoader can be faster if no IO operations when caching all the data in memory
val_loader = ThreadDataLoader(val_ds,
num_workers=0,
batch_size=args.batch_size,
shuffle=False)
# create network, loss function and optimizer
if args.network == "SegResNet":
model = SegResNet(
blocks_down=[1, 2, 2, 4],
blocks_up=[1, 1, 1],
init_filters=16,
in_channels=4,
out_channels=3,
dropout_prob=0.0,
).to(device)
else:
model = UNet(
spatial_dims=3,
in_channels=4,
out_channels=3,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
).to(device)
loss_function = DiceFocalLoss(
smooth_nr=1e-5,
smooth_dr=1e-5,
squared_pred=True,
to_onehot_y=False,
sigmoid=True,
batch=True,
)
optimizer = Novograd(model.parameters(), lr=args.lr)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=args.epochs)
# wrap the model with DistributedDataParallel module
model = DistributedDataParallel(model, device_ids=[device])
dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True,
reduction="mean_batch")
post_trans = Compose(
[EnsureType(),
Activations(sigmoid=True),
AsDiscrete(threshold=0.5)])
# start a typical PyTorch training
best_metric = -1
best_metric_epoch = -1
print(f"time elapsed before training: {time.time() - total_start}")
train_start = time.time()
for epoch in range(args.epochs):
epoch_start = time.time()
print("-" * 10)
print(f"epoch {epoch + 1}/{args.epochs}")
epoch_loss = train(train_loader, model, loss_function, optimizer,
lr_scheduler, scaler)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
if (epoch + 1) % args.val_interval == 0:
metric, metric_tc, metric_wt, metric_et = evaluate(
model, val_loader, dice_metric, dice_metric_batch, post_trans)
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
if dist.get_rank() == 0:
torch.save(model.state_dict(), "best_metric_model.pth")
print(
f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
f" tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f}"
f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
print(
f"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}"
)
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch},"
f" total train time: {(time.time() - train_start):.4f}")
dist.destroy_process_group()
def main():
parser = argparse.ArgumentParser(description="training")
parser.add_argument(
"--checkpoint",
type=str,
default=None,
help="checkpoint full path",
)
parser.add_argument(
"--factor_ram_cost",
default=0.0,
type=float,
help="factor to determine RAM cost in the searched architecture",
)
parser.add_argument(
"--fold",
action="store",
required=True,
help="fold index in N-fold cross-validation",
)
parser.add_argument(
"--json",
action="store",
required=True,
help="full path of .json file",
)
parser.add_argument(
"--json_key",
action="store",
required=True,
help="selected key in .json data list",
)
parser.add_argument(
"--local_rank",
required=int,
help="local process rank",
)
parser.add_argument(
"--num_folds",
action="store",
required=True,
help="number of folds in cross-validation",
)
parser.add_argument(
"--output_root",
action="store",
required=True,
help="output root",
)
parser.add_argument(
"--root",
action="store",
required=True,
help="data root",
)
args = parser.parse_args()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
if not os.path.exists(args.output_root):
os.makedirs(args.output_root, exist_ok=True)
amp = True
determ = True
factor_ram_cost = args.factor_ram_cost
fold = int(args.fold)
input_channels = 1
learning_rate = 0.025
learning_rate_arch = 0.001
learning_rate_milestones = np.array([0.4, 0.8])
num_images_per_batch = 1
num_epochs = 1430 # around 20k iteration
num_epochs_per_validation = 100
num_epochs_warmup = 715
num_folds = int(args.num_folds)
num_patches_per_image = 1
num_sw_batch_size = 6
output_classes = 3
overlap_ratio = 0.625
patch_size = (96, 96, 96)
patch_size_valid = (96, 96, 96)
spacing = [1.0, 1.0, 1.0]
print("factor_ram_cost", factor_ram_cost)
# deterministic training
if determ:
set_determinism(seed=0)
# initialize the distributed training process, every GPU runs in a process
dist.init_process_group(backend="nccl", init_method="env://")
# dist.barrier()
world_size = dist.get_world_size()
with open(args.json, "r") as f:
json_data = json.load(f)
split = len(json_data[args.json_key]) // num_folds
list_train = json_data[args.json_key][:(
split * fold)] + json_data[args.json_key][(split * (fold + 1)):]
list_valid = json_data[args.json_key][(split * fold):(split * (fold + 1))]
# training data
files = []
for _i in range(len(list_train)):
str_img = os.path.join(args.root, list_train[_i]["image"])
str_seg = os.path.join(args.root, list_train[_i]["label"])
if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
continue
files.append({"image": str_img, "label": str_seg})
train_files = files
random.shuffle(train_files)
train_files_w = train_files[:len(train_files) // 2]
train_files_w = partition_dataset(data=train_files_w,
shuffle=True,
num_partitions=world_size,
even_divisible=True)[dist.get_rank()]
print("train_files_w:", len(train_files_w))
train_files_a = train_files[len(train_files) // 2:]
train_files_a = partition_dataset(data=train_files_a,
shuffle=True,
num_partitions=world_size,
even_divisible=True)[dist.get_rank()]
print("train_files_a:", len(train_files_a))
# validation data
files = []
for _i in range(len(list_valid)):
str_img = os.path.join(args.root, list_valid[_i]["image"])
str_seg = os.path.join(args.root, list_valid[_i]["label"])
if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
continue
files.append({"image": str_img, "label": str_seg})
val_files = files
val_files = partition_dataset(data=val_files,
shuffle=False,
num_partitions=world_size,
even_divisible=False)[dist.get_rank()]
print("val_files:", len(val_files))
# network architecture
device = torch.device(f"cuda:{args.local_rank}")
torch.cuda.set_device(device)
train_transforms = Compose([
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(keys=["image", "label"],
pixdim=spacing,
mode=("bilinear", "nearest"),
align_corners=(True, True)),
CastToTyped(keys=["image"], dtype=(torch.float32)),
ScaleIntensityRanged(keys=["image"],
a_min=-87.0,
a_max=199.0,
b_min=0.0,
b_max=1.0,
clip=True),
CastToTyped(keys=["image", "label"], dtype=(np.float16, np.uint8)),
CopyItemsd(keys=["label"], times=1, names=["label4crop"]),
Lambdad(
keys=["label4crop"],
func=lambda x: np.concatenate(tuple([
ndimage.binary_dilation(
(x == _k).astype(x.dtype), iterations=48).astype(x.dtype)
for _k in range(output_classes)
]),
axis=0),
overwrite=True,
),
EnsureTyped(keys=["image", "label"]),
CastToTyped(keys=["image"], dtype=(torch.float32)),
SpatialPadd(keys=["image", "label", "label4crop"],
spatial_size=patch_size,
mode=["reflect", "constant", "constant"]),
RandCropByLabelClassesd(keys=["image", "label"],
label_key="label4crop",
num_classes=output_classes,
ratios=[
1,
] * output_classes,
spatial_size=patch_size,
num_samples=num_patches_per_image),
Lambdad(keys=["label4crop"], func=lambda x: 0),
RandRotated(keys=["image", "label"],
range_x=0.3,
range_y=0.3,
range_z=0.3,
mode=["bilinear", "nearest"],
prob=0.2),
RandZoomd(keys=["image", "label"],
min_zoom=0.8,
max_zoom=1.2,
mode=["trilinear", "nearest"],
align_corners=[True, None],
prob=0.16),
RandGaussianSmoothd(keys=["image"],
sigma_x=(0.5, 1.15),
sigma_y=(0.5, 1.15),
sigma_z=(0.5, 1.15),
prob=0.15),
RandScaleIntensityd(keys=["image"], factors=0.3, prob=0.5),
RandShiftIntensityd(keys=["image"], offsets=0.1, prob=0.5),
RandGaussianNoised(keys=["image"], std=0.01, prob=0.15),
RandFlipd(keys=["image", "label"], spatial_axis=0, prob=0.5),
RandFlipd(keys=["image", "label"], spatial_axis=1, prob=0.5),
RandFlipd(keys=["image", "label"], spatial_axis=2, prob=0.5),
CastToTyped(keys=["image", "label"],
dtype=(torch.float32, torch.uint8)),
ToTensord(keys=["image", "label"]),
])
val_transforms = Compose([
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(keys=["image", "label"],
pixdim=spacing,
mode=("bilinear", "nearest"),
align_corners=(True, True)),
CastToTyped(keys=["image"], dtype=(torch.float32)),
ScaleIntensityRanged(keys=["image"],
a_min=-87.0,
a_max=199.0,
b_min=0.0,
b_max=1.0,
clip=True),
CastToTyped(keys=["image", "label"], dtype=(np.float32, np.uint8)),
EnsureTyped(keys=["image", "label"]),
ToTensord(keys=["image", "label"])
])
train_ds_a = monai.data.CacheDataset(data=train_files_a,
transform=train_transforms,
cache_rate=1.0,
num_workers=8)
train_ds_w = monai.data.CacheDataset(data=train_files_w,
transform=train_transforms,
cache_rate=1.0,
num_workers=8)
val_ds = monai.data.CacheDataset(data=val_files,
transform=val_transforms,
cache_rate=1.0,
num_workers=2)
# monai.data.Dataset can be used as alternatives when debugging or RAM space is limited.
# train_ds_a = monai.data.Dataset(data=train_files_a, transform=train_transforms)
# train_ds_w = monai.data.Dataset(data=train_files_w, transform=train_transforms)
# val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
train_loader_a = ThreadDataLoader(train_ds_a,
num_workers=0,
batch_size=num_images_per_batch,
shuffle=True)
train_loader_w = ThreadDataLoader(train_ds_w,
num_workers=0,
batch_size=num_images_per_batch,
shuffle=True)
val_loader = ThreadDataLoader(val_ds,
num_workers=0,
batch_size=1,
shuffle=False)
# DataLoader can be used as alternatives when ThreadDataLoader is less efficient.
# train_loader_a = DataLoader(train_ds_a, batch_size=num_images_per_batch, shuffle=True, num_workers=2, pin_memory=torch.cuda.is_available())
# train_loader_w = DataLoader(train_ds_w, batch_size=num_images_per_batch, shuffle=True, num_workers=2, pin_memory=torch.cuda.is_available())
# val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=2, pin_memory=torch.cuda.is_available())
dints_space = monai.networks.nets.TopologySearch(
channel_mul=0.5,
num_blocks=12,
num_depths=4,
use_downsample=True,
device=device,
)
model = monai.networks.nets.DiNTS(
dints_space=dints_space,
in_channels=input_channels,
num_classes=output_classes,
use_downsample=True,
)
model = model.to(device)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
post_pred = Compose(
[EnsureType(),
AsDiscrete(argmax=True, to_onehot=output_classes)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=output_classes)])
# loss function
loss_func = monai.losses.DiceCELoss(
include_background=False,
to_onehot_y=True,
softmax=True,
squared_pred=True,
batch=True,
smooth_nr=0.00001,
smooth_dr=0.00001,
)
# optimizer
optimizer = torch.optim.SGD(model.weight_parameters(),
lr=learning_rate * world_size,
momentum=0.9,
weight_decay=0.00004)
arch_optimizer_a = torch.optim.Adam([dints_space.log_alpha_a],
lr=learning_rate_arch * world_size,
betas=(0.5, 0.999),
weight_decay=0.0)
arch_optimizer_c = torch.optim.Adam([dints_space.log_alpha_c],
lr=learning_rate_arch * world_size,
betas=(0.5, 0.999),
weight_decay=0.0)
print()
if torch.cuda.device_count() > 1:
if dist.get_rank() == 0:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = DistributedDataParallel(model,
device_ids=[device],
find_unused_parameters=True)
if args.checkpoint != None and os.path.isfile(args.checkpoint):
print("[info] fine-tuning pre-trained checkpoint {0:s}".format(
args.checkpoint))
model.load_state_dict(torch.load(args.checkpoint, map_location=device))
torch.cuda.empty_cache()
else:
print("[info] training from scratch")
# amp
if amp:
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
if dist.get_rank() == 0:
print("[info] amp enabled")
# start a typical PyTorch training
val_interval = num_epochs_per_validation
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
idx_iter = 0
metric_values = list()
if dist.get_rank() == 0:
writer = SummaryWriter(
log_dir=os.path.join(args.output_root, "Events"))
with open(os.path.join(args.output_root, "accuracy_history.csv"),
"a") as f:
f.write("epoch\tmetric\tloss\tlr\ttime\titer\n")
dataloader_a_iterator = iter(train_loader_a)
start_time = time.time()
for epoch in range(num_epochs):
decay = 0.5**np.sum([
(epoch - num_epochs_warmup) /
(num_epochs - num_epochs_warmup) > learning_rate_milestones
])
lr = learning_rate * decay
for param_group in optimizer.param_groups:
param_group["lr"] = lr
if dist.get_rank() == 0:
print("-" * 10)
print(f"epoch {epoch + 1}/{num_epochs}")
print("learning rate is set to {}".format(lr))
model.train()
epoch_loss = 0
loss_torch = torch.zeros(2, dtype=torch.float, device=device)
epoch_loss_arch = 0
loss_torch_arch = torch.zeros(2, dtype=torch.float, device=device)
step = 0
for batch_data in train_loader_w:
step += 1
inputs, labels = batch_data["image"].to(
device), batch_data["label"].to(device)
if world_size == 1:
for _ in model.weight_parameters():
_.requires_grad = True
else:
for _ in model.module.weight_parameters():
_.requires_grad = True
dints_space.log_alpha_a.requires_grad = False
dints_space.log_alpha_c.requires_grad = False
optimizer.zero_grad()
if amp:
with autocast():
outputs = model(inputs)
if output_classes == 2:
loss = loss_func(torch.flip(outputs, dims=[1]),
1 - labels)
else:
loss = loss_func(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
outputs = model(inputs)
if output_classes == 2:
loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
else:
loss = loss_func(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
loss_torch[0] += loss.item()
loss_torch[1] += 1.0
epoch_len = len(train_loader_w)
idx_iter += 1
if dist.get_rank() == 0:
print("[{0}] ".format(str(datetime.now())[:19]) +
f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")
writer.add_scalar("train_loss", loss.item(),
epoch_len * epoch + step)
if epoch < num_epochs_warmup:
continue
try:
sample_a = next(dataloader_a_iterator)
except StopIteration:
dataloader_a_iterator = iter(train_loader_a)
sample_a = next(dataloader_a_iterator)
inputs_search, labels_search = sample_a["image"].to(
device), sample_a["label"].to(device)
if world_size == 1:
for _ in model.weight_parameters():
_.requires_grad = False
else:
for _ in model.module.weight_parameters():
_.requires_grad = False
dints_space.log_alpha_a.requires_grad = True
dints_space.log_alpha_c.requires_grad = True
# linear increase topology and RAM loss
entropy_alpha_c = torch.tensor(0.).to(device)
entropy_alpha_a = torch.tensor(0.).to(device)
ram_cost_full = torch.tensor(0.).to(device)
ram_cost_usage = torch.tensor(0.).to(device)
ram_cost_loss = torch.tensor(0.).to(device)
topology_loss = torch.tensor(0.).to(device)
probs_a, arch_code_prob_a = dints_space.get_prob_a(child=True)
entropy_alpha_a = -((probs_a) * torch.log(probs_a + 1e-5)).mean()
entropy_alpha_c = -(F.softmax(dints_space.log_alpha_c, dim=-1) * \
F.log_softmax(dints_space.log_alpha_c, dim=-1)).mean()
topology_loss = dints_space.get_topology_entropy(probs_a)
ram_cost_full = dints_space.get_ram_cost_usage(inputs.shape,
full=True)
ram_cost_usage = dints_space.get_ram_cost_usage(inputs.shape)
ram_cost_loss = torch.abs(factor_ram_cost -
ram_cost_usage / ram_cost_full)
arch_optimizer_a.zero_grad()
arch_optimizer_c.zero_grad()
combination_weights = (epoch - num_epochs_warmup) / (
num_epochs - num_epochs_warmup)
if amp:
with autocast():
outputs_search = model(inputs_search)
if output_classes == 2:
loss = loss_func(torch.flip(outputs_search, dims=[1]),
1 - labels_search)
else:
loss = loss_func(outputs_search, labels_search)
loss += combination_weights * ((entropy_alpha_a + entropy_alpha_c) + ram_cost_loss \
+ 0.001 * topology_loss)
scaler.scale(loss).backward()
scaler.step(arch_optimizer_a)
scaler.step(arch_optimizer_c)
scaler.update()
else:
outputs_search = model(inputs_search)
if output_classes == 2:
loss = loss_func(torch.flip(outputs_search, dims=[1]),
1 - labels_search)
else:
loss = loss_func(outputs_search, labels_search)
loss += 1.0 * (combination_weights * (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss \
+ 0.001 * topology_loss)
loss.backward()
arch_optimizer_a.step()
arch_optimizer_c.step()
epoch_loss_arch += loss.item()
loss_torch_arch[0] += loss.item()
loss_torch_arch[1] += 1.0
if dist.get_rank() == 0:
print(
"[{0}] ".format(str(datetime.now())[:19]) +
f"{step}/{epoch_len}, train_loss_arch: {loss.item():.4f}")
writer.add_scalar("train_loss_arch", loss.item(),
epoch_len * epoch + step)
# synchronizes all processes and reduce results
dist.barrier()
dist.all_reduce(loss_torch, op=torch.distributed.ReduceOp.SUM)
loss_torch = loss_torch.tolist()
loss_torch_arch = loss_torch_arch.tolist()
if dist.get_rank() == 0:
loss_torch_epoch = loss_torch[0] / loss_torch[1]
print(
f"epoch {epoch + 1} average loss: {loss_torch_epoch:.4f}, best mean dice: {best_metric:.4f} at epoch {best_metric_epoch}"
)
if epoch >= num_epochs_warmup:
loss_torch_arch_epoch = loss_torch_arch[0] / loss_torch_arch[1]
print(
f"epoch {epoch + 1} average arch loss: {loss_torch_arch_epoch:.4f}, best mean dice: {best_metric:.4f} at epoch {best_metric_epoch}"
)
if (epoch + 1) % val_interval == 0:
torch.cuda.empty_cache()
model.eval()
with torch.no_grad():
metric = torch.zeros((output_classes - 1) * 2,
dtype=torch.float,
device=device)
metric_sum = 0.0
metric_count = 0
metric_mat = []
val_images = None
val_labels = None
val_outputs = None
_index = 0
for val_data in val_loader:
val_images = val_data["image"].to(device)
val_labels = val_data["label"].to(device)
roi_size = patch_size_valid
sw_batch_size = num_sw_batch_size
if amp:
with torch.cuda.amp.autocast():
pred = sliding_window_inference(
val_images,
roi_size,
sw_batch_size,
lambda x: model(x),
mode="gaussian",
overlap=overlap_ratio,
)
else:
pred = sliding_window_inference(
val_images,
roi_size,
sw_batch_size,
lambda x: model(x),
mode="gaussian",
overlap=overlap_ratio,
)
val_outputs = pred
val_outputs = post_pred(val_outputs[0, ...])
val_outputs = val_outputs[None, ...]
val_labels = post_label(val_labels[0, ...])
val_labels = val_labels[None, ...]
value = compute_meandice(y_pred=val_outputs,
y=val_labels,
include_background=False)
print(_index + 1, "/", len(val_loader), value)
metric_count += len(value)
metric_sum += value.sum().item()
metric_vals = value.cpu().numpy()
if len(metric_mat) == 0:
metric_mat = metric_vals
else:
metric_mat = np.concatenate((metric_mat, metric_vals),
axis=0)
for _c in range(output_classes - 1):
val0 = torch.nan_to_num(value[0, _c], nan=0.0)
val1 = 1.0 - torch.isnan(value[0, 0]).float()
metric[2 * _c] += val0 * val1
metric[2 * _c + 1] += val1
_index += 1
# synchronizes all processes and reduce results
dist.barrier()
dist.all_reduce(metric, op=torch.distributed.ReduceOp.SUM)
metric = metric.tolist()
if dist.get_rank() == 0:
for _c in range(output_classes - 1):
print(
"evaluation metric - class {0:d}:".format(_c + 1),
metric[2 * _c] / metric[2 * _c + 1])
avg_metric = 0
for _c in range(output_classes - 1):
avg_metric += metric[2 * _c] / metric[2 * _c + 1]
avg_metric = avg_metric / float(output_classes - 1)
print("avg_metric", avg_metric)
if avg_metric > best_metric:
best_metric = avg_metric
best_metric_epoch = epoch + 1
best_metric_iterations = idx_iter
node_a_d, arch_code_a_d, arch_code_c_d, arch_code_a_max_d = dints_space.decode(
)
torch.save(
{
"node_a": node_a_d,
"arch_code_a": arch_code_a_d,
"arch_code_a_max": arch_code_a_max_d,
"arch_code_c": arch_code_c_d,
"iter_num": idx_iter,
"epochs": epoch + 1,
"best_dsc": best_metric,
"best_path": best_metric_iterations,
},
os.path.join(args.output_root,
"search_code_" + str(idx_iter) + ".pth"),
)
print("saved new best metric model")
dict_file = {}
dict_file["best_avg_dice_score"] = float(best_metric)
dict_file["best_avg_dice_score_epoch"] = int(
best_metric_epoch)
dict_file["best_avg_dice_score_iteration"] = int(idx_iter)
with open(os.path.join(args.output_root, "progress.yaml"),
"w") as out_file:
documents = yaml.dump(dict_file, stream=out_file)
print(
"current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}"
.format(epoch + 1, avg_metric, best_metric,
best_metric_epoch))
current_time = time.time()
elapsed_time = (current_time - start_time) / 60.0
with open(
os.path.join(args.output_root,
"accuracy_history.csv"), "a") as f:
f.write(
"{0:d}\t{1:.5f}\t{2:.5f}\t{3:.5f}\t{4:.1f}\t{5:d}\n"
.format(epoch + 1, avg_metric, loss_torch_epoch,
lr, elapsed_time, idx_iter))
dist.barrier()
torch.cuda.empty_cache()
print(
f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
if dist.get_rank() == 0:
writer.close()
dist.destroy_process_group()
return
SpatialPadd(keys=("input", "mask"), spatial_size=(1120, 1120)),
RandFlipd(keys=("input", "mask"), prob=0.5, spatial_axis=1),
RandAffined(
keys=("input", "mask"),
prob=0.5,
rotate_range=np.pi / 14.4,
translate_range=(70, 70),
scale_range=(0.1, 0.1),
as_tensor_output=False,
),
RandSpatialCropd(
keys=("input", "mask"),
roi_size=(cfg.img_size[0], cfg.img_size[1]),
random_size=False,
),
RandScaleIntensityd(keys="input", factors=(-0.2, 0.2), prob=0.5),
RandShiftIntensityd(keys="input", offsets=(-51, 51), prob=0.5),
RandLambdad(keys="input", func=lambda x: 255 - x, prob=0.5),
RandCoarseDropoutd(
keys=("input", "mask"),
holes=8,
spatial_size=(1, 1),
max_spatial_size=(102, 102),
prob=0.5,
),
CastToTyped(keys="input", dtype=np.float32),
NormalizeIntensityd(keys="input", nonzero=False),
Lambdad(keys="input", func=lambda x: x.clip(-20, 20)),
EnsureTyped(keys=("input", "mask")),
])
# load 4 Nifti images and stack them together
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Spacingd(
keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest"),
),
Orientationd(keys=["image", "label"], axcodes="RAS"),
RandSpatialCropd(keys=["image", "label"],
roi_size=[128, 128, 64],
random_size=False),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
ToTensord(keys=["image", "label"]),
])
val_transform = Compose([
LoadImaged(keys=["image", "label"]),
AsChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Spacingd(
keys=["image", "label"],
pixdim=(1.5, 1.5, 2.0),
mode=("bilinear", "nearest"),
),
Orientationd(keys=["image", "label"], axcodes="RAS"),
CenterSpatialCropd(keys=["image", "label"], roi_size=[128, 128, 64]),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),