import importlib
import sys
import torch
from src.models.unet3d import unet3d
from torchvision import transforms
from src.dataset.train_val_split import train_val_split
from src.losses.ce_dice_loss import CrossEntropyDiceLoss3D
from src.losses import dice_loss, region_based_loss, new_losses
from src.models.io_model import load_model
from src.train.trainer import Trainer, TrainerArgs
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from src.config import BratsConfiguration
from src.dataset.augmentations import color_augmentations, spatial_augmentations
from src.dataset.utils import dataset, visualization as visualization
from src.models.vnet import vnet, asymm_vnet
from src.logging_conf import logger
from src.dataset.loaders.brats_dataset import BratsDataset
def num_params(net_params):
n_params = sum([p.data.nelement() for p in net_params])
logger.info(f"Number of params: {n_params}")
######## PARAMS
logger.info("Processing Parameters...")
config = BratsConfiguration(sys.argv[1])
model_config = config.get_model_config()
dataset_config = config.get_dataset_config()
basic_config = config.get_basic_config()
patch_size = config.patch_size
tensorboard_logdir = basic_config.get("tensorboard_logs")
checkpoint_path = model_config.get("checkpoint")
batch_size = dataset_config.getint("batch_size")
n_patches = dataset_config.getint("n_patches")
n_classes = dataset_config.getint("classes")
loss = model_config.get("loss")
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
logger.info(f"Device: {device}")
######## DATASET
logger.info("Creating Dataset...")
data, _ = dataset.read_brats(dataset_config.get("train_csv"), lgg_only=dataset_config.getboolean("lgg_only"))
data_train, data_val = train_val_split(data, val_size=0.2)
data_train = data_train * n_patches
data_val = data_val * n_patches
n_modalities = dataset_config.getint("n_modalities") # like color channels
sampling_method = importlib.import_module(dataset_config.get("sampling_method"))
transform = transforms.Compose([color_augmentations.RandomIntensityShift(),
color_augmentations.RandomIntensityScale(),
spatial_augmentations.RandomMirrorFlip(p=0.5),
spatial_augmentations.RandomRotation90(p=0.5)])
compute_patch = basic_config.getboolean("compute_patches")
train_dataset = BratsDataset(data_train, sampling_method, patch_size, compute_patch=compute_patch, transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
val_dataset = BratsDataset(data_val, sampling_method, patch_size, compute_patch=compute_patch, transform=transform)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
if basic_config.getboolean("plot"):
data_batch, labels_batch = next(iter(train_loader))
data_batch.reshape(data_batch.shape[0] * data_batch.shape[1], data_batch.shape[2], data_batch.shape[3],
data_batch.shape[4], data_batch.shape[5])
labels_batch.reshape(labels_batch.shape[0] * labels_batch.shape[1], labels_batch.shape[2], labels_batch.shape[3],
labels_batch.shape[4])
print(data_batch.shape)
logger.info('Plotting images')
visualization.plot_batch_slice(data_batch, labels_batch, slice=30, save=True)
######## MODEL
logger.info("Initiating Model...")
config_network = model_config["network"]
if config_network== "vnet":
network = vnet.VNet(elu=model_config.getboolean("use_elu"),
in_channels=n_modalities,
classes=n_classes,
init_features_maps=model_config.getint("init_features_maps"))
elif config_network == "vnet_asymm":
network = asymm_vnet.VNet(non_linearity=model_config.get("non_linearity"), in_channels=n_modalities, classes=n_classes,
init_features_maps=model_config.getint("init_features_maps"), kernel_size=model_config.getint("kernel_size"),
padding=model_config.getint("padding"))
elif config_network == "3dunet_residual":
network = unet3d.ResidualUNet3D(in_channels=n_modalities, out_channels=n_classes, final_sigmoid=False,
f_maps=model_config.getint("init_features_maps"), layer_order="crg",
num_levels=4, num_groups=4,conv_padding=1)
elif config_network == "3dunet":
network = unet3d.UNet3D(in_channels=n_modalities, out_channels=n_classes, final_sigmoid=False,
f_maps=model_config.getint("init_features_maps"), layer_order="crg",
num_levels=4, num_groups=4,conv_padding=1)
else:
raise ValueError("Bad parameter for network {}".format(model_config.get("network")))
num_params(network.parameters())
##### TRAIN
logger.info("Start Training")
network.to(device)
optim = model_config.get("optimizer")
if optim == "SGD":
optimizer = torch.optim.SGD(network.parameters(), lr=model_config.getfloat("learning_rate"),
momentum=model_config.getfloat("momentum"), weight_decay=model_config.getfloat("weight_decay"))
elif optim == "ADAM":
optimizer = torch.optim.Adam(network.parameters(), lr=model_config.getfloat("learning_rate"), weight_decay=model_config.getfloat("weight_decay"), amsgrad=False)
else:
raise ValueError("Bad optimizer. Current options: [SGD, ADAM]")
best_loss = 1000
if basic_config.getboolean("resume"):
logger.info("Loading model from checkpoint..")
model, optimizer, start_epoch, best_loss = load_model(network, checkpoint_path, device, optimizer, True)
logger.info(f"Loaded model with starting epoch {start_epoch}")
else:
start_epoch = 0
writer = SummaryWriter(tensorboard_logdir)
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=model_config.getfloat("lr_decay"),
patience=model_config.getint("patience"))
if loss == "dice":
criterion = dice_loss.DiceLoss(classes=n_classes, eval_regions=model_config.getboolean("eval_regions"),
sigmoid_normalization=True)
elif loss == "combined":
# 0. back, 1: ncr, 2: ed, 3: et
ce_weigh = torch.tensor([0.1, 0.35, 0.2 , 0.35])
criterion = CrossEntropyDiceLoss3D(weight=ce_weigh, classes=n_classes,
eval_regions=model_config.getboolean("eval_regions"), sigmoid_normalization=True)
elif loss == "both_dice":
criterion = region_based_loss.RegionBasedDiceLoss3D(classes=n_classes, sigmoid_normalization=True)
elif loss == "gdl":
criterion = new_losses.GeneralizedDiceLoss()
else:
raise ValueError(f"Bad loss value {loss}. Expected ['dice', combined]")
args = TrainerArgs(model_config.getint("n_epochs"), device, model_config.get("model_path"), loss)
trainer = Trainer(args, network, optimizer, criterion, start_epoch, train_loader, val_loader, scheduler, writer)
trainer.start(best_loss=best_loss)
print("Finished!")
