def __init__(self, model, optimizer, lr_scheduler, loss_criterion, eval_criterion, device, loaders, checkpoint_dir, model_name, max_num_epochs=100, max_num_iterations=1e5, validate_after_iters=100, log_after_iters=100, validate_iters=None, num_iterations=0, num_epoch=0, eval_score_higher_is_better=True, best_eval_score=None, logger=None): # if logger is None: # self.logger = utils.get_logger('VaeUnetTrainer', level=logging.DEBUG) # else: # self.logger = logger self.logger = logger self.config = load_config() self.logger.info(model) self.model = model self.optimizer = optimizer self.scheduler = lr_scheduler self.loss_criterion = loss_criterion self.eval_criterion = eval_criterion self.device = 0 self.loaders = loaders self.checkpoint_dir = checkpoint_dir self.model_name = model_name self.max_num_epochs = max_num_epochs self.max_num_iterations = max_num_iterations self.validate_after_iters = validate_after_iters self.log_after_iters = log_after_iters self.validate_iters = validate_iters self.eval_score_higher_is_better = eval_score_higher_is_better self.use_graph_brain = config['loaders']['graph_brain'] logger.info( f'eval_score_higher_is_better: {eval_score_higher_is_better}') if best_eval_score is not None: self.best_eval_score = best_eval_score else: # initialize the best_eval_score if eval_score_higher_is_better: self.best_eval_score = float('-inf') else: self.best_eval_score = float('+inf') self.writer = SummaryWriter( logdir=os.path.join(checkpoint_dir, self._get_job_name())) self.num_iterations = num_iterations self.num_epoch = num_epoch
def main(): config = load_config() # make sure those values correspond to the ones used during training in_channels = config['in_channels'] out_channels = config['out_channels'] # use the same upsampling as during training interpolate = config['interpolate'] # specify the layer ordering used during training layer_order = config['layer_order'] # should sigmoid be used as a final activation layer final_sigmoid = config['final_sigmoid'] init_channel_number = config['init_channel_number'] model = UNet3D(in_channels, out_channels, init_channel_number=init_channel_number, final_sigmoid=final_sigmoid, interpolate=interpolate, conv_layer_order=layer_order) model_path = config['model_path'] logger.info(f'Loading model from {model_path}...') utils.load_checkpoint(model_path, model) device = config['device'] model = model.to(device) logger.info('Loading datasets...') for test_dataset in get_test_datasets(config): # run the model prediction on the entire dataset probability_maps = predict(model, test_dataset, out_channels, device) # save the resulting probability maps output_file = _get_output_file(test_dataset) save_predictions(probability_maps, output_file)
def main(): # Load configuration config = load_config() # create logger logfile = config.get('logfile', None) logger = utils.get_logger('UNet3DPredictor', logfile=logfile) # Create the model model = get_model(config) # multiple GPUs if (torch.cuda.device_count() > 1): logger.info("There are {} GPUs available".format( torch.cuda.device_count())) model = nn.DataParallel(model) # Load model state model_path = config['model_path'] logger.info(f'Loading model from {model_path}...') utils.load_checkpoint(model_path, model) logger.info(f"Sending the model to '{config['device']}'") model = model.to(config['device']) logger.info('Loading HDF5 datasets...') for test_loader in get_test_loaders(config): logger.info(f"Processing '{test_loader.dataset.file_path}'...") #output_file = _get_output_file(test_loader.dataset) output_file = _get_output_file(config['output_folder'], test_loader.dataset) logger.info(output_file) predictor = _get_predictor(model, test_loader, output_file, config) # run the model prediction on the entire dataset and save to the 'output_file' H5 predictor.predict()
def main(): # Create main logger logger = get_logger('UNet3DTrainer') # Load and log experiment configuration config = load_config() logger.info(config) manual_seed = config.get('manual_seed', None) if manual_seed is not None: logger.info(f'Seed the RNG for all devices with {manual_seed}') torch.manual_seed(manual_seed) # see https://pytorch.org/docs/stable/notes/randomness.html torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # Create the model model = get_model(config) # put the model on GPUs logger.info(f"Sending the model to '{config['device']}'") model = model.to(config['device']) # Log the number of learnable parameters logger.info( f'Number of learnable params {get_number_of_learnable_parameters(model)}' ) # Create loss criterion loss_criterion = get_loss_criterion(config) # Create evaluation metric eval_criterion = get_evaluation_metric(config) # Cross validation path_to_folder = config['loaders']['all_data_path'][0] cross_walidation = CrossValidation(path_to_folder, 1, 3, 2) train_set = cross_walidation.train_filepaths val_set = cross_walidation.validation_filepaths config['loaders']['train_path'] = train_set config['loaders']['val_path'] = val_set # Create data loaders loaders = get_train_loaders(config) # Create the optimizer optimizer = _create_optimizer(config, model) # Create learning rate adjustment strategy lr_scheduler = _create_lr_scheduler(config, optimizer) # Create model trainer trainer = _create_trainer(config, model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, loss_criterion=loss_criterion, eval_criterion=eval_criterion, loaders=loaders, logger=logger) # Start training trainer.fit()
def main(): # Load configuration config = load_config() # Create the model model = get_model(config) # Load model state model_path = config['model_path'] logger.info(f'Loading model from {model_path}...') utils.load_checkpoint(model_path, model) logger.info(f"Sending the model to '{config['device']}'") model = model.to(config['device']) logger.info('Loading HDF5 datasets...') store_predictions_in_memory = config.get('store_predictions_in_memory', True) if store_predictions_in_memory: logger.info( 'Predictions will be stored in memory. Make sure you have enough RAM for you dataset.' ) for test_loader in get_test_loaders(config): logger.info(f"Processing '{test_loader.dataset.file_path}'...") output_file = _get_output_file(test_loader.dataset) # run the model prediction on the entire dataset and save to the 'output_file' H5 if store_predictions_in_memory: predict_in_memory(model, test_loader, output_file, config) else: predict(model, test_loader, output_file, config)
def main(): logger = get_logger('UNet3DTrainer') config = load_config() logger.info(config) manual_seed = config.get('manual_seed', None) if manual_seed is not None: logger.info(f'Seed the RNG for all devices with {manual_seed}') torch.manual_seed(manual_seed) # see https://pytorch.org/docs/stable/notes/randomness.html torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False #loaders = get_train_loaders(config) #print(loaders['test'].__len__()) #for i, t in enumerate(loaders['test']): # for tt in t: # print(i, tt.shape) for loader in get_test_loaders(config): print(loader.dataset.__getid__()) print(loader.dataset.__len__())
def main(): # Load and log experiment configuration config = load_config() # Create main logger logger = get_logger('UNet3DTrainer', file_name=config['trainer']['checkpoint_dir']) logger.info(config) os.environ['CUDA_VISIBLE_DEVICES'] = config['default_device'] assert torch.cuda.is_available(), "Currently, we only support CUDA version" manual_seed = config.get('manual_seed', None) if manual_seed is not None: logger.info(f'Seed the RNG for all devices with {manual_seed}') torch.manual_seed(manual_seed) # see https://pytorch.org/docs/stable/notes/randomness.html torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # Create the model model = get_model(config) # model, parameters = generate_model(MedConfig) # put the model on GPUs logger.info(f"Sending the model to '{config['default_device']}'") model = torch.nn.DataParallel(model).cuda() # Log the number of learnable parameters logger.info( f'Number of learnable params {get_number_of_learnable_parameters(model)}' ) # Create loss criterion loss_criterion = get_loss_criterion(config) # Create evaluation metric eval_criterion = get_evaluation_metric(config) # Create data loaders loaders = get_brats_train_loaders(config) # Create the optimizer optimizer = _create_optimizer(config, model) # Create learning rate adjustment strategy lr_scheduler = _create_lr_scheduler(config, optimizer) # Create model trainer trainer = _create_trainer(config, model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, loss_criterion=loss_criterion, eval_criterion=eval_criterion, loaders=loaders, logger=logger) # Start training trainer.fit()
def main(): # Load and log experiment configuration config = load_config() # Create main logger logfile = config.get('logfile', None) logger = get_logger('UNet3DTrainer', logfile=logfile) logger.info(config) manual_seed = config.get('manual_seed', None) if manual_seed is not None: logger.info(f'Seed the RNG for all devices with {manual_seed}') torch.manual_seed(manual_seed) # see https://pytorch.org/docs/stable/notes/randomness.html torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # Create the model model = get_model(config) # multiple GPUs if (torch.cuda.device_count() > 1): logger.info("There are {} GPUs available".format(torch.cuda.device_count())) model = nn.DataParallel(model) # put the model on GPUs logger.info(f"Sending the model to '{config['device']}'") model = model.to(config['device']) # Log the number of learnable parameters logger.info(f'Number of learnable params {get_number_of_learnable_parameters(model)}') # Create loss criterion loss_criterion = get_loss_criterion(config) logger.info(f"Created loss criterion: {config['loss']['name']}") # Create evaluation metric eval_criterion = get_evaluation_metric(config) logger.info(f"Created eval criterion: {config['eval_metric']['name']}") # Create data loaders loaders = get_train_loaders(config) # Create the optimizer optimizer = _create_optimizer(config, model) # Create learning rate adjustment strategy lr_scheduler = _create_lr_scheduler(config, optimizer) # Create model trainer trainer = _create_trainer(config, model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, loss_criterion=loss_criterion, eval_criterion=eval_criterion, loaders=loaders, logger=logger) # Start training trainer.fit()
def main(): # Load and log experiment configuration config = load_config() logger.info(config) # exit() manual_seed = config.get('manual_seed', None) if manual_seed is not None: logger.info(f'Seed the RNG for all devices with {manual_seed}') torch.manual_seed(manual_seed) # see https://pytorch.org/docs/stable/notes/randomness.html torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # Create the model model = get_model(config) # use DataParallel if more than 1 GPU available device = config['device'] if torch.cuda.device_count() > 1 and not device.type == 'cpu': model = nn.DataParallel(model) logger.info(f'Using {torch.cuda.device_count()} GPUs for training') # put the model on GPUs logger.info(f"Sending the model to '{config['device']}'") model = model.to(device) # Log the number of learnable parameters logger.info( f'Number of learnable params {get_number_of_learnable_parameters(model)}' ) # Create loss criterion loss_criterion = get_loss_criterion(config) # Create evaluation metric eval_criterion = get_evaluation_metric(config) # Create data loaders loaders = get_train_loaders_1(config) # Create the optimizer optimizer = _create_optimizer(config, model) # Create learning rate adjustment strategy lr_scheduler = _create_lr_scheduler(config, optimizer) # Create model trainer trainer = _create_trainer(config, model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, loss_criterion=loss_criterion, eval_criterion=eval_criterion, loaders=loaders) # Start training trainer.fit()
def main(): logger = get_logger('UNet3DTrainer') config = load_config() logger.info(config) # Create loss criterion loss_criterion = get_loss_criterion(config) # Create the model model = UNet3D(config['in_channels'], config['out_channels'], final_sigmoid=config['final_sigmoid'], init_channel_number=config['init_channel_number'], conv_layer_order=config['layer_order'], interpolate=config['interpolate']) model = model.to(config['device']) # Log the number of learnable parameters logger.info(f'Number of learnable params {get_number_of_learnable_parameters(model)}') # Create evaluation metric eval_criterion = get_evaluation_metric(config) loaders = get_train_loaders(config) # Create the optimizer optimizer = _create_optimizer(config, model) # Create learning rate adjustment strategy lr_scheduler = _create_lr_scheduler(config, optimizer) if config['resume'] is not None: trainer = UNet3DTrainer.from_checkpoint(config['resume'], model, optimizer, lr_scheduler, loss_criterion, eval_criterion, loaders, logger=logger) else: trainer = UNet3DTrainer(model, optimizer, lr_scheduler, loss_criterion, eval_criterion, config['device'], loaders, config['checkpoint_dir'], max_num_epochs=config['epochs'], max_num_iterations=config['iters'], validate_after_iters=config['validate_after_iters'], log_after_iters=config['log_after_iters'], logger=logger) trainer.fit()
def main(): # Create main logger logger = get_logger('UNet3DTrainer') # Load and log experiment configuration config = load_config() # Set DEFAULT_DEVICE and config file logger.info(config) # Log configure from train_config_4d_input.yaml manual_seed = config.get('manual_seed', None) if manual_seed is not None: logger.info(f'Seed the RNG for all devices with {manual_seed}') torch.manual_seed(manual_seed) torch.backends.cudnn.deterministic = True # Ensure the repeatability of the experiment torch.backends.cudnn.benchmark = False # Benchmark mode improves the computation speed, but results in slightly different network feedforward results # Create the model model = get_model(config) # put the model on GPUs logger.info(f"Sending the model to '{config['device']}'") model = model.to(config['device']) # Log the number of learnable parameters logger.info(f'Number of learnable params {get_number_of_learnable_parameters(model)}') # Create loss criterion loss_criterion = get_loss_criterion(config) # Create evaluation metric eval_criterion = get_evaluation_metric(config) # Create data loaders # loaders: {'train': train_loader, 'val': val_loader} loaders = get_train_loaders(config) # Create the optimizer optimizer = _create_optimizer(config, model) # Create learning rate adjustment strategy lr_scheduler = _create_lr_scheduler(config, optimizer) # Create model trainer trainer = _create_trainer(config, model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, loss_criterion=loss_criterion, eval_criterion=eval_criterion, loaders=loaders, logger=logger) # Start training trainer.fit()
def main(): # Load configuration config = load_config() # Create the model model = get_model(config) # Load model state model_path = config['model_path'] logger.info(f'Loading model from {model_path}...') utils.load_checkpoint(model_path, model) logger.info(f"Sending the model to '{config['device']}'") model = model.to(config['device']) logger.info('Loading HDF5 datasets...') for test_loader in get_test_loaders(config): logger.info(f"Processing '{test_loader.dataset.file_path}'...") output_file = _get_output_file(test_loader.dataset) # run the model prediction on the entire dataset and save to the 'output_file' H5 predict(model, test_loader, output_file, config)
def main(): # Load configuration config = load_config() # Create the model model = get_model(config) # Load model state model_path = config['model_path'] logger.info(f'Loading model from {model_path}...') utils.load_checkpoint(model_path, model) logger.info(f"Sending the model to '{config['device']}'") model = model.to(config['device']) logger.info('Loading HDF5 datasets...') test_loader = get_test_loaders(config)['test'] for i, data_pair in enumerate(test_loader): output_file = 'predict_' + str(i) + '.h5' predictor = _get_predictor(model, data_pair, output_file, config) predictor.predict()
def main(): # Load configuration config = load_config() # Create the model model = get_model(config) # Load model state model_path = config['model_path'] logger.info(f'Loading model from {model_path}...') utils.load_checkpoint(model_path, model) model = model.to(config['device']) logger.info('Loading HDF5 datasets...') for test_dataset in get_test_datasets(config): logger.info(f"Processing '{test_dataset.file_path}'...") # run the model prediction on the entire dataset predictions = predict(model, test_dataset, config) # save the resulting probability maps output_file = _get_output_file(test_dataset) dataset_names = _get_dataset_names(config, len(predictions)) save_predictions(predictions, output_file, dataset_names)
def main(): # Load configuration config = load_config() # Create the model model = get_model(config) # Create evaluation metric eval_criterion = get_evaluation_metric(config) # Load model state model_path = config['model_path'] logger.info(f'Loading model from {model_path}...') utils.load_checkpoint(model_path, model) model = model.to(config['device']) logger.info('Loading HDF5 datasets...') # ========================== for data batch, score recording ========================== nii_path="/data/cephfs/punim0877/liver_segmentation_v1/Test_Batch" # load path of test batch hdf5_path="./resources/hdf5ed_test_data" # create dir to save predict image stage=1 for index in range(110,131): # delete for loop. only need one file if not hdf5_it(nii_path,hdf5_path,index,stage): continue config["datasets"]["test_path"]=[] for hdf5_file in os.listdir(hdf5_path): print("adding %s to trainging list" % (hdf5_file)) config["datasets"]["test_path"].append(os.path.join(hdf5_path,hdf5_file)) for test_dataset in get_test_datasets(config): logger.info(f"Processing '{test_dataset.file_path}'...") # run the model prediction on the entire dataset predictions = predict(model, test_dataset, config, eval_criterion) # save the resulting probability maps output_file = _get_output_file(test_dataset) dataset_names = _get_dataset_names(config, len(predictions)) save_predictions(predictions, output_file, dataset_names)
import numpy as np import torch import datetime import cv2 import matplotlib.pyplot as plt plt.switch_backend('agg') from unet3d.config import load_config from tensorboardX import SummaryWriter from torch.optim.lr_scheduler import ReduceLROnPlateau from tqdm import tqdm from visualization import board_add_images, board_add_image from . import utils config = load_config() class UNet3DTrainer: """3D UNet trainer. Args: model (Unet3D): UNet 3D model to be trained optimizer (nn.optim.Optimizer): optimizer used for training lr_scheduler (torch.optim.lr_scheduler._LRScheduler): learning rate scheduler WARN: bear in mind that lr_scheduler.step() is invoked after every validation step (i.e. validate_after_iters) not after every epoch. So e.g. if one uses StepLR with step_size=30 the learning rate will be adjusted after every 30 * validate_after_iters iterations. loss_criterion (callable): loss function eval_criterion (callable): used to compute training/validation metric (such as Dice, IoU, AP or Rand score) saving the best checkpoint is based on the result of this function on the validation set device (torch.device): device to train on