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
