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 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(): # 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 _train_save_load(self, tmpdir, loss, val_metric, model='UNet3D', max_num_epochs=1, log_after_iters=2, validate_after_iters=2, max_num_iterations=4, weight_map=False): binary_loss = loss in ['BCEWithLogitsLoss', 'DiceLoss', 'GeneralizedDiceLoss'] device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu') test_config = copy.deepcopy(CONFIG_BASE) test_config['model']['name'] = model test_config.update({ # get device to train on 'device': device, 'loss': {'name': loss, 'weight': np.random.rand(2).astype(np.float32)}, 'eval_metric': {'name': val_metric} }) test_config['model']['final_sigmoid'] = binary_loss if weight_map: test_config['loaders']['weight_internal_path'] = 'weight_map' loss_criterion = get_loss_criterion(test_config) eval_criterion = get_evaluation_metric(test_config) model = get_model(test_config) model = model.to(device) if loss in ['BCEWithLogitsLoss']: label_dtype = 'float32' else: label_dtype = 'long' test_config['loaders']['transformer']['train']['label'][0]['dtype'] = label_dtype test_config['loaders']['transformer']['test']['label'][0]['dtype'] = label_dtype train, val = TestUNet3DTrainer._create_random_dataset((128, 128, 128), (64, 64, 64), binary_loss) test_config['loaders']['train_path'] = [train] test_config['loaders']['val_path'] = [val] loaders = get_train_loaders(test_config) optimizer = _create_optimizer(test_config, model) test_config['lr_scheduler']['name'] = 'MultiStepLR' lr_scheduler = _create_lr_scheduler(test_config, optimizer) logger = get_logger('UNet3DTrainer', logging.DEBUG) formatter = DefaultTensorboardFormatter() trainer = UNet3DTrainer(model, optimizer, lr_scheduler, loss_criterion, eval_criterion, device, loaders, tmpdir, max_num_epochs=max_num_epochs, log_after_iters=log_after_iters, validate_after_iters=validate_after_iters, max_num_iterations=max_num_iterations, logger=logger, tensorboard_formatter=formatter) trainer.fit() # test loading the trainer from the checkpoint trainer = UNet3DTrainer.from_checkpoint(os.path.join(tmpdir, 'last_checkpoint.pytorch'), model, optimizer, lr_scheduler, loss_criterion, eval_criterion, loaders, logger=logger, tensorboard_formatter=formatter) return trainer
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(): # 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)
from tensorboardX import SummaryWriter from visualization import board_add_images, board_add_image def get_job_name(): now = '{:%Y-%m-%d.%H:%M}'.format(datetime.datetime.now()) return "%s_model" % (now) logger = utils.get_logger('UNet3DPredictor') # Load and log experiment configuration config = load_config() # Load model state model = get_model(config) model_path = config['trainer']['test_model'] logger.info(f'Loading model from {model_path}...') utils.load_checkpoint(model_path, model) # Run on GPU or CPU # if torch.cuda.is_available(): # print("using cuda (", torch.cuda.device_count(), "device(s))") # if torch.cuda.device_count() > 1: # model = nn.DataParallel(model) # device = torch.device("cuda:1") # else: # device = torch.device("cpu") # print("using cpu") # model = model.to(device) logger.info(f"Sending the model to '{config['device']}'")
import torch from unet3d import utils from unet3d import config from unet3d.model import get_model import numpy as np import os in_channels = 1 out_channels = 2 final_sigmoid = False config_file_path = os.environ['UNET_CONFIG_PATH'] config = config._load_config_yaml(config_file_path) InstantiatedModel = get_model(config) # InstantiatedModel = model.UNet3D( in_channels, # out_channels, # final_sigmoid, # f_maps=32, # layer_order='crg', # num_groups=8) InstantiatedModel.training = False def pre_process(input_numpy_patch): input_numpy_patch *= 255 # chunkflow scales integer values to [0,1] input_numpy_patch = (input_numpy_patch - 124.7) / 54.5 #img = np.squeeze(input_numpy_patch)
if __name__ == '__main__': # Load configuration config = load_config() # Load model state model_path = config['model_path'] model_fd = Path(model_path).parent logger = call_logger(log_file=str(model_fd / 'test_log.txt'), log_name='UNetPredict') # Create the model model = get_model(config, is_test=True) if 'output_path' in config.keys(): out_path = config['output_path'] else: out_path = str(model_fd / 'h5_pred') os.makedirs(out_path, exist_ok=True) 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...') datasets_config = config['datasets']