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']
