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():
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 _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():
parser = _arg_parser()
logger = get_logger('UNet3DTrainer')
# Get device to train on
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
args = parser.parse_args()
logger.info(args)
# Create loss criterion
loss_criterion, final_sigmoid = _get_loss_criterion(args.loss)
model = _create_model(args.in_channels,
args.out_channels,
layer_order=args.layer_order,
interpolate=args.interpolate,
final_sigmoid=final_sigmoid)
model = model.to(device)
# Log the number of learnable parameters
logger.info(
f'Number of learnable params {get_number_of_learnable_parameters(model)}'
)
# Create error criterion
error_criterion = DiceCoefficient()
# Get data loaders
loaders = _get_loaders(args.config_dir, logger)
# Create the optimizer
optimizer = _create_optimizer(args, model)
if args.resume:
trainer = UNet3DTrainer.from_checkpoint(
args.resume,
model,
optimizer,
loss_criterion,
error_criterion,
loaders,
validate_after_iters=args.validate_after_iters,
log_after_iters=args.log_after_iters,
logger=logger)
else:
trainer = UNet3DTrainer(model,
optimizer,
loss_criterion,
error_criterion,
device,
loaders,
args.checkpoint_dir,
validate_after_iters=args.validate_after_iters,
log_after_iters=args.log_after_iters,
logger=logger)
trainer.fit()
def __init__(self, model, loader, output_file, config, **kwargs):
self.model = model
self.loader = loader
self.output_file = output_file
self.config = config
self.predictor_config = kwargs
self.logfile = self.config.get('logfile', None)
self.logger = get_logger('UNet3DTrainer', logfile=self.logfile)
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 _train_save_load(self,
tmpdir,
loss,
max_num_epochs=1,
log_after_iters=2,
validate_after_iters=2,
max_num_iterations=4):
# get device to train on
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
# conv-relu-groupnorm
conv_layer_order = 'crg'
final_sigmoid = loss == 'bce'
loss_criterion = get_loss_criterion(loss,
final_sigmoid,
weight=torch.rand(2).to(device))
model = self._create_model(final_sigmoid, conv_layer_order)
accuracy_criterion = DiceCoefficient()
channel_per_class = loss == 'bce'
if loss in ['bce', 'dice']:
label_dtype = 'float32'
else:
label_dtype = 'long'
pixel_wise_weight = loss == 'pce'
loaders = self._get_loaders(channel_per_class=channel_per_class,
label_dtype=label_dtype,
pixel_wise_weight=pixel_wise_weight)
learning_rate = 2e-4
weight_decay = 0.0001
optimizer = optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
logger = get_logger('UNet3DTrainer', logging.DEBUG)
trainer = UNet3DTrainer(model,
optimizer,
loss_criterion,
accuracy_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)
trainer.fit()
# test loading the trainer from the checkpoint
trainer = UNet3DTrainer.from_checkpoint(os.path.join(
tmpdir, 'last_checkpoint.pytorch'),
model,
optimizer,
loss_criterion,
accuracy_criterion,
loaders,
logger=logger)
return trainer
def get_brats_train_loaders(config):
"""
Returns dictionary containing the training and validation loaders
(torch.utils.data.DataLoader) backed by the datasets.hdf5.HDF5Dataset.
:param config: a top level configuration object containing the 'loaders' key
:return: dict {
'train': <train_loader>
'val': <val_loader>
}
"""
assert 'loaders' in config, 'Could not find data loaders configuration'
loaders_config = config['loaders']
logger = get_logger('BraTS_Dataset')
logger.info('Creating training and validation set loaders...')
# get train and validation files
data_paths = loaders_config['dataset_path']
assert isinstance(data_paths, list)
brats = BraTS.DataSet(brats_root=data_paths[0], year=2019).train
train_ids, test_ids, validation_ids = get_all_partition_ids()
# loss_file_num = 0
# for i in train_ids:
# name = i + ".tfrecord.gzip"
# answer = search('/home/server/data/TFRecord/val', name)
# if answer == -1:
# print("查无此文件", name)
# loss_file_num += 1
# print(f'loss file num is {loss_file_num}')
logger.info(f'Loading training set from: {data_paths}...')
train_datasets = BraTSDataset(brats, train_ids)
logger.info(f'Loading validation set from: {data_paths}...')
brats = BraTS.DataSet(brats_root=data_paths[0], year=2019).train
val_datasets = BraTSDataset(brats, validation_ids)
num_workers = loaders_config.get('num_workers', 1)
logger.info(f'Number of workers for train/val datasets: {num_workers}')
# when training with volumetric data use batch_size of 1 due to GPU memory constraints
return {
'train':
DataLoader(train_datasets,
batch_size=1,
shuffle=True,
num_workers=num_workers),
'val':
DataLoader(val_datasets,
batch_size=1,
shuffle=True,
num_workers=num_workers)
}
def get_test_loaders(config):
"""
Returns a list of DataLoader, one per each test file.
:param config: a top level configuration object containing the 'datasets' key
:return: generator of DataLoader objects
"""
assert 'datasets' in config, 'Could not find data sets configuration'
datasets_config = config['datasets']
logger = get_logger('HDF5Dataset', logfile=config['logfile'])
# get train and validation files
test_paths = datasets_config['test_path']
#print(test_paths)
assert isinstance(test_paths, list)
# get h5 internal path
raw_internal_path = datasets_config['raw_internal_path']
# get train/validation patch size and stride
patch = tuple(datasets_config['patch'])
stride = tuple(datasets_config['stride'])
mirror_padding = datasets_config.get('mirror_padding', False)
pad_width = datasets_config.get('pad_width', 20)
if mirror_padding:
logger.info(f'Using mirror padding. Pad width: {pad_width}')
num_workers = datasets_config.get('num_workers', 1)
logger.info(f'Number of workers for the dataloader: {num_workers}')
batch_size = datasets_config.get('batch_size', 1)
logger.info(f'Batch size for dataloader: {batch_size}')
# construct datasets lazily
datasets = (HDF5Dataset(test_path,
patch,
stride,
phase='test',
raw_internal_path=raw_internal_path,
transformer_config=datasets_config['transformer'],
mirror_padding=mirror_padding,
pad_width=pad_width) for test_path in test_paths)
# use generator in order to create data loaders lazily one by one
for dataset in datasets:
logger.info(f'Loading test set from: {dataset.file_path}...')
yield DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
collate_fn=prediction_collate)
def get_test_loaders(config):
"""
Returns a list of DataLoader, one per each test file.
:param config: a top level configuration object containing the 'datasets' key
:return: generator of DataLoader objects
"""
def my_collate(batch):
error_msg = "batch must contain tensors or slice; found {}"
if isinstance(batch[0], torch.Tensor):
return torch.stack(batch, 0)
elif isinstance(batch[0], slice):
return batch[0]
elif isinstance(batch[0], collections.Sequence):
transposed = zip(*batch)
return [my_collate(samples) for samples in transposed]
raise TypeError((error_msg.format(type(batch[0]))))
logger = get_logger('HDF5Dataset')
assert 'datasets' in config, 'Could not find data sets configuration'
datasets_config = config['datasets']
# get train and validation files
test_paths = datasets_config['test_path']
assert isinstance(test_paths, list)
# get h5 internal path
raw_internal_path = datasets_config['raw_internal_path']
# get train/validation patch size and stride
patch = tuple(datasets_config['patch'])
stride = tuple(datasets_config['stride'])
num_workers = datasets_config.get('num_workers', 1)
# construct datasets lazily
datasets = (HDF5Dataset(test_path,
patch,
stride,
phase='test',
raw_internal_path=raw_internal_path,
transformer_config=datasets_config['transformer'])
for test_path in test_paths)
# use generator in order to create data loaders lazily one by one
for dataset in datasets:
logger.info(f'Loading test set from: {dataset.file_path}...')
yield DataLoader(dataset,
batch_size=1,
num_workers=num_workers,
collate_fn=my_collate)
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 test_single_epoch(self, tmpdir, capsys):
with capsys.disabled():
# get device to train on
device = torch.device(
"cuda:0" if torch.cuda.is_available() else 'cpu')
conv_layer_order = 'crg'
loss_criterion, final_sigmoid = DiceLoss(), True
model = self._load_model(final_sigmoid, conv_layer_order)
error_criterion = DiceCoefficient()
loaders = self._get_loaders()
learning_rate = 1e-4
weight_decay = 0.0005
optimizer = optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
logger = get_logger('UNet3DTrainer', logging.DEBUG)
trainer = UNet3DTrainer(model,
optimizer,
loss_criterion,
error_criterion,
device,
loaders,
tmpdir,
max_num_epochs=1,
log_after_iters=2,
validate_after_iters=2,
logger=logger)
trainer.fit()
# test loading the trainer from the checkpoint
UNet3DTrainer.from_checkpoint(os.path.join(
tmpdir, 'last_checkpoint.pytorch'),
model,
optimizer,
loss_criterion,
error_criterion,
loaders,
logger=logger)
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 _train_save_load(self, tmpdir, loss, val_metric, max_num_epochs=1, log_after_iters=2, validate_after_iters=2,
max_num_iterations=4):
# get device to train on
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
# conv-relu-groupnorm
conv_layer_order = 'crg'
final_sigmoid = loss == 'bce'
loss_criterion = get_loss_criterion(loss, weight=torch.rand(2).to(device))
eval_criterion = get_evaluation_metric(val_metric)
model = self._create_model(final_sigmoid, conv_layer_order)
channel_per_class = loss == 'bce'
if loss in ['bce']:
label_dtype = 'float32'
else:
label_dtype = 'long'
pixel_wise_weight = loss == 'pce'
patch = (32, 64, 64)
stride = (32, 64, 64)
train, val = TestUNet3DTrainer._create_random_dataset((128, 128, 128), (64, 64, 64), channel_per_class)
loaders = get_loaders([train], [val], 'raw', 'label', label_dtype=label_dtype, train_patch=patch,
train_stride=stride, val_patch=patch, val_stride=stride, transformer='BaseTransformer',
pixel_wise_weight=pixel_wise_weight)
learning_rate = 2e-4
weight_decay = 0.0001
optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
logger = get_logger('UNet3DTrainer', logging.DEBUG)
trainer = UNet3DTrainer(model, optimizer, 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)
trainer.fit()
# test loading the trainer from the checkpoint
trainer = UNet3DTrainer.from_checkpoint(
os.path.join(tmpdir, 'last_checkpoint.pytorch'),
model, optimizer, loss_criterion, eval_criterion, loaders,
logger=logger)
return trainer
import importlib
import os
import numpy as np
import time
import torch
import torch.nn.functional as F
from skimage import measure
import hdbscan
from sklearn.cluster import MeanShift
from unet3d.losses import compute_per_channel_dice
from unet3d.utils import get_logger, adapted_rand, expand_as_one_hot, plot_segm
LOGGER = get_logger('EvalMetric')
class DiceCoefficient:
"""Computes Dice Coefficient.
Generalized to multiple channels by computing per-channel Dice Score
(as described in https://arxiv.org/pdf/1707.03237.pdf) and theTn simply taking the average.
Input is expected to be probabilities instead of logits.
This metric is mostly useful when channels contain the same semantic class (e.g. affinities computed with different offsets).
DO NOT USE this metric when training with DiceLoss, otherwise the results will be biased towards the loss.
"""
def __init__(self, epsilon=1e-5, ignore_index=None, **kwargs):
self.epsilon = epsilon
self.ignore_index = ignore_index
def __call__(self, input, target):
"""
def __init__(self,
file_path,
patch_shape,
stride_shape,
phase,
transformer_config,
raw_internal_path='raw',
label_internal_path='label',
weight_internal_path=None,
slice_builder_cls=SliceBuilder,
mirror_padding=False,
pad_width=20,
logfile=None):
"""
:param file_path: path to H5 file containing raw data as well as labels and per pixel weights (optional)
:param patch_shape: the shape of the patch DxHxW
:param stride_shape: the shape of the stride DxHxW
:param phase: 'train' for training, 'val' for validation, 'test' for testing; data augmentation is performed
only during the 'train' phase
:param transformer_config: data augmentation configuration
:param raw_internal_path (str or list): H5 internal path to the raw dataset
:param label_internal_path (str or list): H5 internal path to the label dataset
:param weight_internal_path (str or list): H5 internal path to the per pixel weights
:param slice_builder_cls: defines how to sample the patches from the volume
:param mirror_padding (bool): pad with the reflection of the vector mirrored on the first and last values
along each axis. Only applicable during the 'test' phase
:param pad_width: number of voxels padded to the edges of each axis (only if `mirror_padding=True`)
"""
assert phase in ['train', 'val', 'test']
self._check_patch_shape(patch_shape)
self.phase = phase
self.file_path = file_path
self.mirror_padding = mirror_padding
self.pad_width = pad_width
self.logfile = logfile
self.logger = get_logger('HDF5Dataset', logfile=self.logfile)
# convert raw_internal_path, label_internal_path and weight_internal_path to list for ease of computation
if isinstance(raw_internal_path, str):
raw_internal_path = [raw_internal_path]
if isinstance(label_internal_path, str):
label_internal_path = [label_internal_path]
if isinstance(weight_internal_path, str):
weight_internal_path = [weight_internal_path]
with h5py.File(file_path, 'r') as input_file:
# WARN: we load everything into memory due to hdf5 bug when reading H5 from multiple subprocesses, i.e.
# File "h5py/_proxy.pyx", line 84, in h5py._proxy.H5PY_H5Dread
# OSError: Can't read data (inflate() failed)
self.raws = [
input_file[internal_path][...]
for internal_path in raw_internal_path
]
# calculate global mean and std for Normalization augmentation
mean, std = self._calculate_mean_std(self.raws[0])
self.transformer = transforms.get_transformer(
transformer_config, mean, std, phase)
self.raw_transform = self.transformer.raw_transform()
if phase != 'test':
# create label/weight transform only in train/val phase
self.label_transform = self.transformer.label_transform()
self.labels = [
input_file[internal_path][...]
for internal_path in label_internal_path
]
if weight_internal_path is not None:
# look for the weight map in the raw file
self.weight_maps = [
input_file[internal_path][...]
for internal_path in weight_internal_path
]
self.weight_transform = self.transformer.weight_transform()
else:
self.weight_maps = None
self._check_dimensionality(self.raws, self.labels)
else:
# 'test' phase used only for predictions so ignore the label dataset
self.labels = None
self.weight_maps = None
# add mirror padding if needed
if self.mirror_padding:
padded_volumes = []
for raw in self.raws:
if raw.ndim == 4:
channels = [
np.pad(r,
pad_width=self.pad_width,
mode='reflect') for r in raw
]
padded_volume = np.stack(channels)
else:
padded_volume = np.pad(raw,
pad_width=self.pad_width,
mode='reflect')
padded_volumes.append(padded_volume)
self.raws = padded_volumes
# build slice indices for raw and label data sets
slice_builder = slice_builder_cls(self.raws, self.labels,
self.weight_maps, patch_shape,
stride_shape)
self.raw_slices = slice_builder.raw_slices
self.label_slices = slice_builder.label_slices
self.weight_slices = slice_builder.weight_slices
self.patch_count = len(self.raw_slices)
self.logger.info(f'Number of patches: {self.patch_count}')
import os
import h5py
import numpy as np
import torch
from datasets.hdf5 import get_test_datasets
from unet3d import utils
from unet3d.config import load_config
from unet3d.model import get_model
logger = utils.get_logger('UNet3DPredictor')
def predict(model, hdf5_dataset, config):
"""
Return prediction masks by applying the model on the given dataset
Args:
model (Unet3D): trained 3D UNet model used for prediction
hdf5_dataset (torch.utils.data.Dataset): input dataset
out_channels (int): number of channels in the network output
device (torch.Device): device to run the prediction on
Returns:
prediction_maps (numpy array): prediction masks for given dataset
"""
def _volume_shape(hdf5_dataset):
#TODO: support multiple internal datasets
raw = hdf5_dataset.raws[0]
import hdbscan
import numpy as np
import torch
from skimage import measure
from skimage.metrics import adapted_rand_error, peak_signal_noise_ratio
from sklearn.cluster import MeanShift
# from pytorch3dunet.unet3d.losses import compute_per_channel_dice
# from pytorch3dunet.unet3d.seg_metrics import AveragePrecision, Accuracy
# from pytorch3dunet.unet3d.utils import get_logger, expand_as_one_hot, plot_segm, convert_to_numpy
from unet3d.losses import compute_per_channel_dice
from unet3d.seg_metrics import AveragePrecision, Accuracy
from unet3d.utils import get_logger, expand_as_one_hot, plot_segm, convert_to_numpy
logger = get_logger('EvalMetric')
class DiceCoefficient:
"""Computes Dice Coefficient.
Generalized to multiple channels by computing per-channel Dice Score
(as described in https://arxiv.org/pdf/1707.03237.pdf) and theTn simply taking the average.
Input is expected to be probabilities instead of logits.
This metric is mostly useful when channels contain the same semantic class (e.g. affinities computed with different offsets).
DO NOT USE this metric when training with DiceLoss, otherwise the results will be biased towards the loss.
"""
def __init__(self, epsilon=1e-6, **kwargs):
self.epsilon = epsilon
def __call__(self, input, target):
# Average across channels in order to get the final score
import collections
import importlib
import h5py
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, ConcatDataset
import augment.transforms as transforms
from unet3d.utils import get_logger
logger = get_logger('HDF5Dataset')
class SliceBuilder:
def __init__(self, raw_datasets, label_datasets, weight_dataset, patch_shape, stride_shape):
self._raw_slices = self._build_slices(raw_datasets[0], patch_shape, stride_shape)
if label_datasets is None:
self._label_slices = None
else:
# take the first element in the label_datasets to build slices
self._label_slices = self._build_slices(label_datasets[0], patch_shape, stride_shape)
assert len(self._raw_slices) == len(self._label_slices)
if weight_dataset is None:
self._weight_slices = None
else:
self._weight_slices = self._build_slices(weight_dataset[0], patch_shape, stride_shape)
assert len(self.raw_slices) == len(self._weight_slices)
@property
def raw_slices(self):
def get_test_loaders(config):
"""
Returns a list of DataLoader, one per each test file.
:param config: a top level configuration object containing the 'datasets' key
:return: generator of DataLoader objects
"""
assert 'datasets' in config, 'Could not find data sets configuration'
datasets_config = config['datasets']
logfile = config.get('logfile', None)
logger = get_logger('CloudVolumeDataset', logfile=logfile)
# get test data information
image_cv_path = datasets_config['image_cv_path']
seg_cv_path = datasets_config['seg_cv_path']
#cutout_pkl_file = datasets_config['cutout_pkl_file']
mip_level = datasets_config['mip_level']
volume_start = datasets_config['volume_start']
volume_end = datasets_config['volume_end']
id = datasets_config['id']
cutout_bounds = [
volume_start[0], volume_end[0], volume_start[1], volume_end[1],
volume_start[2], volume_end[2]
]
image_cv = cloudvolume.CloudVolume(image_cv_path,
mip=mip_level,
bounded=False,
autocrop=True,
fill_missing=True)
seg_cv = None
if not (seg_cv_path is None):
seg_cv = cloudvolume.CloudVolume(seg_cv_path,
mip=mip_level,
bounded=False,
autocrop=True,
fill_missing=True)
patch = tuple(datasets_config['patch'])
stride = tuple(datasets_config['stride'])
mirror_padding = datasets_config.get('mirror_padding', False)
pad_width = datasets_config.get('pad_width', 20)
if mirror_padding:
logger.info(f'Using mirror padding. Pad width: {pad_width}')
num_workers = datasets_config.get('num_workers', 1)
logger.info(f'Number of workers for the dataloader: {num_workers}')
batch_size = datasets_config.get('batch_size', 1)
logger.info(f'Batch size for dataloader: {batch_size}')
#with open(cutout_pkl_file, 'rb') as f:
# df = pickle.load(f)
test_datasets = []
test_dataset = CloudVolumeDataset(
image_cv,
seg_cv,
id,
cutout_bounds,
mip_level,
'test',
patch,
stride,
transformer_config=datasets_config['transformer'],
mirror_padding=mirror_padding,
pad_width=pad_width)
test_datasets.append(test_dataset)
#for indx, row in df.iterrows():
# if row['id'] == 0 or row['phase'] != 'test':
# continue
# test_dataset = CloudVolumeDataset(image_cv, seg_cv, row['id'], list(row['cutout_bounds']),
# mip_level, 'test', patch, stride,
# transformer_config=datasets_config['transformer'],
# mirror_padding=mirror_padding, pad_width=pad_width)
# test_datasets.append(test_dataset)
# use generator in order to create data loaders lazily one by one
for indx, dataset in enumerate(test_datasets):
logger.info(f'Loading test set no: {indx}')
yield DataLoader(dataset,
batch_size=batch_size,
num_workers=num_workers,
collate_fn=prediction_collate)
def __init__(self,
image_cv,
seg_cv,
id,
bounds,
mip_level,
phase,
patch_shape,
stride_shape,
transformer_config,
slice_builder_cls=SliceBuilder,
mirror_padding=False,
pad_width=20,
logfile=None):
assert phase in ['train', 'val', 'test']
self.phase = phase
self._check_patch_shape(patch_shape)
self.image_cv = image_cv
self.seg_cv = seg_cv
self.id = id
assert isinstance(bounds, list)
self.bounds = bounds
self.mip_level = mip_level
self.mirror_padding = mirror_padding
self.pad_width = pad_width
self.logger = get_logger('CloudVolumeDataset', logfile=logfile)
minx, maxx, miny, maxy, minz, maxz = self.bounds
self.raws = []
img = np.squeeze(self.image_cv[minx:maxx, miny:maxy, minz:maxz, 0])
# transpose (the data is always CDHW)
img = np.transpose(img, (2, 0, 1))
self.raws.append(img)
mean, std = self._calculate_mean_std(self.raws[0])
self.transformer = transforms.get_transformer(transformer_config, mean,
std, phase)
self.raw_transform = self.transformer.raw_transform()
if phase != 'test':
# create label/weight transform only in train/val phase
self.label_transform = self.transformer.label_transform()
label = np.squeeze(self.seg_cv[minx:maxx, miny:maxy, minz:maxz, 0])
label = np.where(label / np.ndarray.max(label) >= 0.2, 1, 0)
# transpose the label (the data is always CDHW)
label = np.transpose(label, (2, 0, 1))
self.labels = [label]
self.weight_maps = None
self._check_dimensionality(self.raws, self.labels)
else:
self.labels = None
self.weight_maps = None
# add mirror padding if needed
if self.mirror_padding:
padded_volumes = [
np.pad(raw, pad_width=self.pad_width, mode='reflect')
for raw in self.raws
]
self.raws = padded_volumes
#print(self.raws[0].shape, self.labels[0].shape)
#print(np.min(self.labels[0][:,:,200]), np.max(self.labels[0][:,:,200]))
#plt.imshow(self.labels[0][:,:,200])
#plt.show()
#plt.imshow(self.raws[0][:,:,200])
#plt.show()
# build slice indices for raw and label data sets
slice_builder = slice_builder_cls(self.raws, self.labels,
self.weight_maps, patch_shape,
stride_shape)
self.raw_slices = slice_builder.raw_slices
self.label_slices = slice_builder.label_slices
self.weight_slices = slice_builder.weight_slices
self.patch_count = len(self.raw_slices)
self.logger.info(f'Number of patches: {self.patch_count}')
def get_train_loaders(config):
"""
Returns dictionary containing the training and validation loaders
:param config: a top level configuration object containing the 'loaders' key
:return: dict {
'train': <train_loader>
'val': <val_loader>
}
"""
assert 'loaders' in config, 'Could not find data loaders configuration'
loaders_config = config['loaders']
logger = get_logger('CloudVolumeDataset', logfile=config['logfile'])
logger.info('Creating training and validation set loaders...')
# get image cloudvolume path and segmentation mask cv path
image_cv_path = loaders_config['image_cv_path']
seg_cv_path = loaders_config['seg_cv_path']
cutout_pkl_file = loaders_config['cutout_pkl_file']
mip_level = loaders_config['mip_level']
#volume_start = dataset_config['volume_start']
#volume_end = dataset_config['volume_end']
# get train/validation patch size and stride
train_patch = tuple(loaders_config['train_patch'])
train_stride = tuple(loaders_config['train_stride'])
val_patch = tuple(loaders_config['val_patch'])
val_stride = tuple(loaders_config['val_stride'])
# get train slice_builder_cls
train_slice_builder_str = loaders_config.get('train_slice_builder',
'SliceBuilder')
logger.info(f'Train slice builder class: {train_slice_builder_str}')
train_slice_builder_cls = _get_slice_builder_cls(train_slice_builder_str)
image_cv = cloudvolume.CloudVolume(image_cv_path,
mip=mip_level,
bounded=False,
autocrop=True,
fill_missing=True)
seg_cv = cloudvolume.CloudVolume(seg_cv_path,
mip=mip_level,
bounded=False,
autocrop=True,
fill_missing=True)
df = []
with open(cutout_pkl_file, 'rb') as f:
df = pickle.load(f)
train_datasets = []
val_datasets = []
val_slice_builder_str = loaders_config.get('val_slice_builder',
'SliceBuilder')
logger.info(f'Val slice builder class: {val_slice_builder_str}')
val_slice_builder = _get_slice_builder_cls(val_slice_builder_str)
phase = 'train'
for indx, row in df.iterrows():
if row['id'] == 0 or row['phase'] == 'test':
continue
row['cutout_bounds'] = list(row['cutout_bounds'])
if row['phase'] == 'train':
try:
logger.info(
"Loading image and segmentation mask for {}".format(
row['id']))
train_dataset = CloudVolumeDataset(
image_cv,
seg_cv,
row['id'],
row['cutout_bounds'],
mip_level,
'train',
train_patch,
train_stride,
transformer_config=loaders_config['transformer'],
slice_builder_cls=train_slice_builder_cls)
train_datasets.append(train_dataset)
except Exception:
logger.info("Skipping training data for: {}".format(row['id']),
exc_info=True)
if row['phase'] == 'val':
try:
logger.info(
f"Loading image and segmentation mask for {row['id']}")
val_dataset = CloudVolumeDataset(
image_cv,
seg_cv,
row['id'],
row['cutout_bounds'],
mip_level,
'val',
val_patch,
val_stride,
transformer_config=loaders_config['transformer'],
slice_builder_cls=val_slice_builder)
val_datasets.append(val_dataset)
except Exception:
logger.info(f"Skipping validation data for: {row['id']}",
exc_info=True)
num_workers = loaders_config.get('num_workers', 1)
logger.info(f'Number of workers for train/val dataloader: {num_workers}')
batch_size = loaders_config.get('batch_size', 1)
logger.info(f'Batch size for train/val loader: {batch_size}')
# when training with volumetric data use batch_size of 1 due to GPU memory constraints
return {
'train':
DataLoader(ConcatDataset(train_datasets),
batch_size=batch_size,
shuffle=True,
num_workers=num_workers),
'val':
DataLoader(ConcatDataset(val_datasets),
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
}
def main():
logger = get_logger('UNet3DTrainer')
# Get device to train on
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
config = parse_train_config()
logger.info(config)
# Create loss criterion
if config.loss_weight is not None:
loss_weight = torch.tensor(config.loss_weight)
loss_weight = loss_weight.to(device)
else:
loss_weight = None
loss_criterion = get_loss_criterion(config.loss, loss_weight,
config.ignore_index)
model = UNet3D(config.in_channels,
config.out_channels,
init_channel_number=config.init_channel_number,
conv_layer_order=config.layer_order,
interpolate=config.interpolate,
final_sigmoid=config.final_sigmoid)
model = model.to(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.eval_metric,
ignore_index=config.ignore_index)
# Get data loaders. If 'bce' or 'dice' loss is used, convert labels to float
train_path, val_path = config.train_path, config.val_path
if config.loss in ['bce']:
label_dtype = 'float32'
else:
label_dtype = 'long'
train_patch = tuple(config.train_patch)
train_stride = tuple(config.train_stride)
val_patch = tuple(config.val_patch)
val_stride = tuple(config.val_stride)
logger.info(f'Train patch/stride: {train_patch}/{train_stride}')
logger.info(f'Val patch/stride: {val_patch}/{val_stride}')
pixel_wise_weight = config.loss == 'pce'
loaders = get_loaders(train_path,
val_path,
label_dtype=label_dtype,
raw_internal_path=config.raw_internal_path,
label_internal_path=config.label_internal_path,
train_patch=train_patch,
train_stride=train_stride,
val_patch=val_patch,
val_stride=val_stride,
transformer=config.transformer,
pixel_wise_weight=pixel_wise_weight,
curriculum_learning=config.curriculum,
ignore_index=config.ignore_index)
# Create the optimizer
optimizer = _create_optimizer(config, model)
if config.resume:
trainer = UNet3DTrainer.from_checkpoint(config.resume,
model,
optimizer,
loss_criterion,
eval_criterion,
loaders,
logger=logger)
else:
trainer = UNet3DTrainer(
model,
optimizer,
loss_criterion,
eval_criterion,
device,
loaders,
config.checkpoint_dir,
max_num_epochs=config.epochs,
max_num_iterations=config.iters,
max_patience=config.patience,
validate_after_iters=config.validate_after_iters,
log_after_iters=config.log_after_iters,
logger=logger)
trainer.fit()
def main():
parser = _arg_parser()
logger = get_logger('UNet3DTrainer')
# Get device to train on
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
args = parser.parse_args()
logger.info(args)
# Create loss criterion
if args.loss_weight is not None:
loss_weight = torch.tensor(args.loss_weight)
loss_weight = loss_weight.to(device)
else:
loss_weight = None
loss_criterion, final_sigmoid = _get_loss_criterion(args.loss, loss_weight)
model = _create_model(args.in_channels, args.out_channels,
layer_order=args.layer_order,
interpolate=args.interpolate,
final_sigmoid=final_sigmoid)
model = model.to(device)
# Log the number of learnable parameters
#logger.info(f'Number of learnable params {get_number_of_learnable_parameters(model)}')
# Create accuracy metric
accuracy_criterion = _get_accuracy_criterion(not final_sigmoid)
# Get data loaders. If 'bce' or 'dice' loss is used, convert labels to float
train_path, val_path = args.train_path, args.val_path
if args.loss in ['bce', 'dice']:
label_dtype = 'float32'
else:
label_dtype = 'long'
train_patch = tuple(args.train_patch)
train_stride = tuple(args.train_stride)
val_patch = tuple(args.val_patch)
val_stride = tuple(args.val_stride)
#logger.info(f'Train patch/stride: {train_patch}/{train_stride}')
#logger.info(f'Val patch/stride: {val_patch}/{val_stride}')
loaders = _get_loaders(train_path, val_path, label_dtype=label_dtype, train_patch=train_patch,
train_stride=train_stride, val_patch=val_patch, val_stride=val_stride)
# Create the optimizer
optimizer = _create_optimizer(args, model)
if args.resume:
trainer = UNet3DTrainer.from_checkpoint(args.resume, model,
optimizer, loss_criterion,
accuracy_criterion, loaders,
logger=logger)
else:
trainer = UNet3DTrainer(model, optimizer, loss_criterion,
accuracy_criterion,
device, loaders, args.checkpoint_dir,
max_num_epochs=args.epochs,
max_num_iterations=args.iters,
max_patience=args.patience,
validate_after_iters=args.validate_after_iters,
log_after_iters=args.log_after_iters,
logger=logger)
trainer.fit()
# Authors: # Christian F. Baumgartner ([email protected]) # Lisa M. Koch ([email protected]) # Source: https://git.ee.ethz.ch/baumgach/discriminative_learning_toolbox/blob/master/utils.py import nibabel as nib import numpy as np import os import logging from unet3d.utils import get_logger from skimage import measure, transform logger = get_logger('utils') try: import cv2 except: logger.warning( 'Could not import opencv. Augmentation functions will be unavailable.') else: def rotate_image(img, angle, interp=cv2.INTER_LINEAR): rows, cols = img.shape[:2] rotation_matrix = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1) out = cv2.warpAffine(img, rotation_matrix, (cols, rows), flags=interp, borderMode=cv2.BORDER_REPLICATE) return np.reshape(out, img.shape) def rotate_image_as_onehot(img, angle, nlabels, interp=cv2.INTER_LINEAR):
def get_train_loaders(config):
"""
Returns dictionary containing the training and validation loaders
(torch.utils.data.DataLoader) backed by the datasets.hdf5.HDF5Dataset.
:param config: a top level configuration object containing the 'loaders' key
:return: dict {
'train': <train_loader>
'val': <val_loader>
}
"""
assert 'loaders' in config, 'Could not find data loaders configuration'
loaders_config = config['loaders']
logger = get_logger('HDF5Dataset')
logger.info('Creating training and validation set loaders...')
# get train and validation files
train_paths = loaders_config['train_path']
val_paths = loaders_config['val_path']
assert isinstance(train_paths, list)
assert isinstance(val_paths, list)
# get h5 internal paths for raw and label
raw_internal_path = loaders_config['raw_internal_path']
label_internal_path = loaders_config['label_internal_path']
weight_internal_path = loaders_config.get('weight_internal_path', None)
# get train/validation patch size and stride
train_patch = tuple(loaders_config['train_patch'])
train_stride = tuple(loaders_config['train_stride'])
val_patch = tuple(loaders_config['val_patch'])
val_stride = tuple(loaders_config['val_stride'])
# get slice_builder_cls
slice_builder_str = loaders_config.get('slice_builder', 'SliceBuilder')
logger.info(f'Slice builder class: {slice_builder_str}')
slice_builder_cls = _get_slice_builder_cls(slice_builder_str)
train_datasets = []
for train_path in train_paths:
try:
logger.info(f'Loading training set from: {train_path}...')
# create H5 backed training and validation dataset with data augmentation
train_dataset = HDF5Dataset(
train_path,
train_patch,
train_stride,
phase='train',
transformer_config=loaders_config['transformer'],
raw_internal_path=raw_internal_path,
label_internal_path=label_internal_path,
weight_internal_path=weight_internal_path,
slice_builder_cls=slice_builder_cls)
train_datasets.append(train_dataset)
except Exception:
logger.info(f'Skipping training set: {train_path}', exc_info=True)
val_datasets = []
for val_path in val_paths:
try:
logger.info(f'Loading validation set from: {train_path}...')
val_dataset = HDF5Dataset(
val_path,
val_patch,
val_stride,
phase='val',
transformer_config=loaders_config['transformer'],
raw_internal_path=raw_internal_path,
label_internal_path=label_internal_path,
weight_internal_path=weight_internal_path)
val_datasets.append(val_dataset)
except Exception:
logger.info(f'Skipping validation set: {val_path}', exc_info=True)
num_workers = loaders_config.get('num_workers', 1)
logger.info(f'Number of workers for train/val datasets: {num_workers}')
# when training with volumetric data use batch_size of 1 due to GPU memory constraints
return {
'train':
DataLoader(ConcatDataset(train_datasets),
batch_size=1,
shuffle=True,
num_workers=num_workers),
'val':
DataLoader(ConcatDataset(val_datasets),
batch_size=1,
shuffle=True,
num_workers=num_workers)
}
import os
import torch
import torch.nn as nn
from tensorboardX import SummaryWriter
from torch.optim.lr_scheduler import ReduceLROnPlateau
# from pytorch3dunet.unet3d.utils import get_logger
from unet3d.utils import get_logger
from . import utils
logger = get_logger('UNet3DTrainer')
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
loaders (dict): 'train' and 'val' loaders
checkpoint_dir (string): dir for saving checkpoints and tensorboard logs
max_num_epochs (int): maximum number of epochs
import collections
import importlib
import numpy as np
import torch
from torch.utils.data import DataLoader, ConcatDataset, Dataset
# from pytorch3dunet.unet3d.utils import get_logger
# from pytorch3dunet.datasets import lung
from unet3d.utils import get_logger
from datasets import lung
import monai.transforms as mn_tf
logger = get_logger('Dataset')
class ConfigDataset(Dataset):
def __getitem__(self, index):
raise NotImplementedError
def __len__(self):
raise NotImplementedError
@classmethod
def create_datasets(cls, dataset_config, phase):
"""
Factory method for creating a list of datasets based on the provided config.
Args:
dataset_config (dict): dataset configuration
