def train(train_loader, model, criterion, optimizer, scheduler, num_classes, batch_size, ep_idx, progress_log, device, debug=False):
"""
Train the model and return the metrics of the training epoch
:param train_loader: training data loader
:param model: model to train
:param criterion: loss criterion
:param optimizer: optimizer to use
:param scheduler: learning rate scheduler
:param num_classes: number of classes
:param batch_size: number of samples to process simultaneously
:param ep_idx: epoch index (for hypertrainer log)
:param progress_log: progress log file (for hypertrainer log)
:param device: device used by pytorch (cpu ou cuda)
:return: Updated training loss
"""
model.train()
train_metrics = create_metrics_dict(num_classes)
with tqdm(train_loader, desc=f'Iterating train batches with {device.type}') as _tqdm:
for batch_index, data in enumerate(_tqdm):
progress_log.open('a', buffering=1).write(tsv_line(ep_idx, 'trn', batch_index, len(train_loader), time.time()))
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
train_metrics['loss'].update(loss.item(), batch_size)
if device.type == 'cuda' and debug:
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(OrderedDict(trn_loss=f'{train_metrics["loss"].val:.2f}',
gpu_perc=f'{res.gpu} %',
gpu_RAM=f'{mem.used / (1024 ** 2):.0f}/{mem.total / (1024 ** 2):.0f} MiB',
lr=optimizer.param_groups[0]['lr'],
img=data['sat_img'].numpy().shape[1:],
smpl=data['map_img'].numpy().shape,
bs=batch_size))
loss.backward()
optimizer.step()
scheduler.step()
logging.info(f'Training Loss: {train_metrics["loss"].avg:.4f}')
return train_metrics
def evaluation(eval_loader,
model,
criterion,
num_classes,
batch_size,
task,
ep_idx,
progress_log,
vis_params,
batch_metrics=None,
dataset='val',
device=None,
debug=False):
"""
Evaluate the model and return the updated metrics
:param eval_loader: data loader
:param model: model to evaluate
:param criterion: loss criterion
:param num_classes: number of classes
:param batch_size: number of samples to process simultaneously
:param task: segmentation or classification
:param ep_idx: epoch index (for hypertrainer log)
:param progress_log: progress log file (for hypertrainer log)
:param batch_metrics: (int) Metrics computed every (int) batches. If left blank, will not perform metrics.
:param dataset: (str) 'val or 'tst'
:param device: device used by pytorch (cpu ou cuda)
:return: (dict) eval_metrics
"""
eval_metrics = create_metrics_dict(num_classes)
model.eval()
vis_at_eval = get_key_def('vis_at_evaluation', vis_params['visualization'],
False)
vis_batch_range = get_key_def('vis_batch_range',
vis_params['visualization'], None)
min_vis_batch, max_vis_batch, increment = vis_batch_range
with tqdm(eval_loader,
dynamic_ncols=True,
desc=f'Iterating {dataset} batches with {device.type}') as _tqdm:
for batch_index, data in enumerate(_tqdm):
progress_log.open('a', buffering=1).write(
tsv_line(ep_idx, dataset, batch_index, len(eval_loader),
time.time()))
with torch.no_grad():
inputs = data['sat_img'].to(device)
labels = data['map_img'].to(device)
labels_flatten = flatten_labels(labels)
outputs = model(inputs)
if isinstance(outputs, OrderedDict):
outputs = outputs['out']
if vis_batch_range is not None and vis_at_eval and batch_index in range(
min_vis_batch, max_vis_batch, increment):
vis_path = progress_log.parent.joinpath('visualization')
if ep_idx == 0 and batch_index == min_vis_batch:
tqdm.write(
f'Visualizing on {dataset} outputs for batches in range {vis_batch_range}. All '
f'images will be saved to {vis_path}\n')
vis_from_batch(params,
inputs,
outputs,
batch_index=batch_index,
vis_path=vis_path,
labels=labels,
dataset=dataset,
ep_num=ep_idx + 1)
outputs_flatten = flatten_outputs(outputs, num_classes)
loss = criterion(outputs, labels)
eval_metrics['loss'].update(loss.item(), batch_size)
if (dataset == 'val') and (batch_metrics is not None):
# Compute metrics every n batches. Time consuming.
assert batch_metrics <= len(_tqdm), f"Batch_metrics ({batch_metrics} is smaller than batch size " \
f"{len(_tqdm)}. Metrics in validation loop won't be computed"
if (
batch_index + 1
) % batch_metrics == 0: # +1 to skip val loop at very beginning
a, segmentation = torch.max(outputs_flatten, dim=1)
eval_metrics = report_classification(
segmentation,
labels_flatten,
batch_size,
eval_metrics,
ignore_index=get_key_def("ignore_index",
params["training"], None))
elif dataset == 'tst':
a, segmentation = torch.max(outputs_flatten, dim=1)
eval_metrics = report_classification(
segmentation,
labels_flatten,
batch_size,
eval_metrics,
ignore_index=get_key_def("ignore_index",
params["training"], None))
_tqdm.set_postfix(
OrderedDict(dataset=dataset,
loss=f'{eval_metrics["loss"].avg:.4f}'))
if debug and device.type == 'cuda':
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(
OrderedDict(
device=device,
gpu_perc=f'{res.gpu} %',
gpu_RAM=
f'{mem.used/(1024**2):.0f}/{mem.total/(1024**2):.0f} MiB'
))
print(f"{dataset} Loss: {eval_metrics['loss'].avg}")
if batch_metrics is not None:
print(f"{dataset} precision: {eval_metrics['precision'].avg}")
print(f"{dataset} recall: {eval_metrics['recall'].avg}")
print(f"{dataset} fscore: {eval_metrics['fscore'].avg}")
return eval_metrics
def train(train_loader,
model,
criterion,
optimizer,
scheduler,
num_classes,
batch_size,
task,
ep_idx,
progress_log,
vis_params,
device,
debug=False):
"""
Train the model and return the metrics of the training epoch
:param train_loader: training data loader
:param model: model to train
:param criterion: loss criterion
:param optimizer: optimizer to use
:param scheduler: learning rate scheduler
:param num_classes: number of classes
:param batch_size: number of samples to process simultaneously
:param task: segmentation or classification
:param ep_idx: epoch index (for hypertrainer log)
:param progress_log: progress log file (for hypertrainer log)
:param vis_params: (dict) Parameters found in the yaml config file. Named vis_params because they are only used for
visualization functions.
:param device: device used by pytorch (cpu ou cuda)
:return: Updated training loss
"""
model.train()
train_metrics = create_metrics_dict(num_classes)
vis_at_train = get_key_def('vis_at_train', vis_params['visualization'],
False)
vis_batch_range = get_key_def('vis_batch_range',
vis_params['visualization'], None)
min_vis_batch, max_vis_batch, increment = vis_batch_range
with tqdm(train_loader,
desc=f'Iterating train batches with {device.type}') as _tqdm:
for batch_index, data in enumerate(_tqdm):
progress_log.open('a', buffering=1).write(
tsv_line(ep_idx, 'trn', batch_index, len(train_loader),
time.time()))
inputs = data['sat_img'].to(device)
labels = data['map_img'].to(device)
# forward
optimizer.zero_grad()
outputs = model(inputs)
# added for torchvision models that output an OrderedDict with outputs in 'out' key.
# More info: https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101/
if isinstance(outputs, OrderedDict):
outputs = outputs['out']
if vis_batch_range is not None and vis_at_train and batch_index in range(
min_vis_batch, max_vis_batch, increment):
vis_path = progress_log.parent.joinpath('visualization')
if ep_idx == 0:
tqdm.write(
f'Visualizing on train outputs for batches in range {vis_batch_range}. All images will be saved to {vis_path}\n'
)
vis_from_batch(params,
inputs,
outputs,
batch_index=batch_index,
vis_path=vis_path,
labels=labels,
dataset='trn',
ep_num=ep_idx + 1)
loss = criterion(outputs, labels)
train_metrics['loss'].update(loss.item(), batch_size)
if device.type == 'cuda' and debug:
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(
OrderedDict(
trn_loss=f'{train_metrics["loss"].val:.2f}',
gpu_perc=f'{res.gpu} %',
gpu_RAM=
f'{mem.used / (1024 ** 2):.0f}/{mem.total / (1024 ** 2):.0f} MiB',
lr=optimizer.param_groups[0]['lr'],
img=data['sat_img'].numpy().shape[1:],
smpl=data['map_img'].numpy().shape,
bs=batch_size,
out_vals=np.unique(
outputs[0].argmax(dim=0).detach().cpu().numpy())))
loss.backward()
optimizer.step()
scheduler.step()
print(f'Training Loss: {train_metrics["loss"].avg:.4f}')
return train_metrics
def sem_seg_inference(model, nd_array, overlay, chunk_size, num_classes, device, meta_map=None, metadata=None, output_path=Path(os.getcwd()), index=0, debug=False):
"""Inference on images using semantic segmentation
Args:
model: model to use for inference
nd_array: nd_array
overlay: amount of overlay to apply
num_classes: number of different classes that may be predicted by the model
device: device used by pytorch (cpu ou cuda)
meta_map:
metadata:
output_path: path to save debug files
index: (int) index of array from list of images on which inference is performed
returns a numpy array of the same size (h,w) as the input image, where each value is the predicted output.
"""
# switch to evaluate mode
model.eval()
if len(nd_array.shape) == 3:
h, w, nb = nd_array.shape
# Pad with overlay on left and top and pad with chunk_size on right and bottom
padded_array = np.pad(nd_array, ((overlay, chunk_size), (overlay, chunk_size), (0, 0)), mode='constant')
elif len(nd_array.shape) == 2:
h, w = nd_array.shape
padded_array = np.expand_dims(np.pad(nd_array, ((overlay, chunk_size), (overlay, chunk_size)),
mode='constant'), axis=0)
else:
h = 0
w = 0
padded_array = None
h_padded, w_padded = padded_array.shape[:2]
# Create an empty array of dimensions (c x h x w): num_classes x height of padded array x width of padded array
output_probs = np.empty([num_classes, h_padded, w_padded], dtype=np.float32)
# Create identical 0-filled array without channels dimension to receive counts for number of outputs generated in specific area.
output_counts = np.zeros([output_probs.shape[1], output_probs.shape[2]], dtype=np.int32)
if padded_array.any():
with torch.no_grad():
for row in tqdm(range(overlay, h + chunk_size, chunk_size - overlay), position=1, leave=False,
desc=f'Inferring rows with "{device}"'):
row_start = row - overlay
row_end = row_start + chunk_size
with tqdm(range(overlay, w + chunk_size, chunk_size - overlay), position=2, leave=False, desc='Inferring columns') as _tqdm:
for col in _tqdm:
col_start = col - overlay
col_end = col_start + chunk_size
chunk_input = padded_array[row_start:row_end, col_start:col_end, :]
if meta_map:
chunk_input = MetaSegmentationDataset.append_meta_layers(chunk_input, meta_map, metadata)
inputs = torch.from_numpy(np.float32(np.transpose(chunk_input, (2, 0, 1))))
inputs.unsqueeze_(0) #Add dummy batch dimension
inputs = inputs.to(device)
# forward
outputs = model(inputs)
# torchvision models give output in 'out' key. May cause problems in future versions of torchvision.
if isinstance(outputs, OrderedDict) and 'out' in outputs.keys():
outputs = outputs['out']
if debug:
if index == 0:
tqdm.write(f'(debug mode) Visualizing inferred tiles...')
vis_from_batch(params, inputs, outputs, batch_index=0, vis_path=output_path,
dataset=f'{row_start}_{col_start}_inf', ep_num=index, debug=True)
outputs = F.softmax(outputs, dim=1)
output_counts[row_start:row_end, col_start:col_end] += 1
# Add inference on sub-image to all completed inferences on previous sub-images.
# FIXME: This operation need to be optimized. Using a lot of RAM on large images.
output_probs[:, row_start:row_end, col_start:col_end] += np.squeeze(outputs.cpu().numpy(),
axis=0)
if debug and device.type == 'cuda':
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(OrderedDict(gpu_perc=f'{res.gpu} %',
gpu_RAM=f'{mem.used / (1024 ** 2):.0f}/{mem.total / (1024 ** 2):.0f} MiB',
inp_size=inputs.cpu().numpy().shape,
out_size=outputs.cpu().numpy().shape,
overlay=overlay))
if debug:
output_counts_PIL = Image.fromarray(output_counts.astype(np.uint8), mode='L')
output_counts_PIL.save(output_path.joinpath(f'output_counts.png'))
tqdm.write(f'Dividing array according to output counts...\n')
# Divide array according to output counts. Manages overlap and returns a softmax array as if only one forward pass had been done.
output_mask_raw = np.divide(output_probs, np.maximum(output_counts, 1)) # , 1 is added to overwrite 0 values.
# Resize the output array to the size of the input image and write it
output_mask_raw_cropped = np.moveaxis(output_mask_raw, 0, -1)
output_mask_raw_cropped = output_mask_raw_cropped[overlay:(h + overlay), overlay:(w + overlay), :]
return output_mask_raw_cropped
else:
raise IOError(f"Error classifying image : Image shape of {len(nd_array.shape)} is not recognized")
def evaluation(eval_loader,
model,
criterion,
num_classes,
batch_size,
ep_idx,
progress_log,
scale,
vis_params,
batch_metrics=None,
dataset='val',
device=None,
debug=False):
"""
Evaluate the model and return the updated metrics
:param eval_loader: data loader
:param model: model to evaluate
:param criterion: loss criterion
:param num_classes: number of classes
:param batch_size: number of samples to process simultaneously
:param ep_idx: epoch index (for hypertrainer log)
:param progress_log: progress log file (for hypertrainer log)
:param scale: Scale to which values in sat img have been redefined. Useful during visualization
:param vis_params: (Dict) Parameters useful during visualization
:param batch_metrics: (int) Metrics computed every (int) batches. If left blank, will not perform metrics.
:param dataset: (str) 'val or 'tst'
:param device: device used by pytorch (cpu ou cuda)
:param debug: if True, debug functions will be performed
:return: (dict) eval_metrics
"""
eval_metrics = create_metrics_dict(num_classes)
model.eval()
for batch_index, data in enumerate(tqdm(eval_loader, dynamic_ncols=True, desc=f'Iterating {dataset} '
f'batches with {device.type}')):
progress_log.open('a', buffering=1).write(tsv_line(ep_idx, dataset, batch_index, len(eval_loader), time.time()))
with torch.no_grad():
try: # For HPC when device 0 not available. Error: RuntimeError: CUDA error: invalid device ordinal
inputs = data['sat_img'].to(device)
labels = data['map_img'].to(device)
except RuntimeError:
logging.exception(f'Unable to use device {device}. Trying "cuda:0"')
device = torch.device('cuda')
inputs = data['sat_img'].to(device)
labels = data['map_img'].to(device)
labels_flatten = flatten_labels(labels)
outputs = model(inputs)
if isinstance(outputs, OrderedDict):
outputs = outputs['out']
# vis_batch_range: range of batches to perform visualization on. see README.md for more info.
# vis_at_eval: (bool) if True, will perform visualization at eval time, as long as vis_batch_range is valid
if vis_params['vis_batch_range'] and vis_params['vis_at_eval']:
min_vis_batch, max_vis_batch, increment = vis_params['vis_batch_range']
if batch_index in range(min_vis_batch, max_vis_batch, increment):
vis_path = progress_log.parent.joinpath('visualization')
if ep_idx == 0 and batch_index == min_vis_batch:
logging.info(
f'Visualizing on {dataset} outputs for batches in range {vis_params["vis_batch_range"]} '
f'images will be saved to {vis_path}\n')
vis_from_batch(vis_params, inputs, outputs,
batch_index=batch_index,
vis_path=vis_path,
labels=labels,
dataset=dataset,
ep_num=ep_idx + 1,
scale=scale)
outputs_flatten = flatten_outputs(outputs, num_classes)
loss = criterion(outputs, labels)
eval_metrics['loss'].update(loss.item(), batch_size)
if (dataset == 'val') and (batch_metrics is not None):
# Compute metrics every n batches. Time consuming.
if not batch_metrics <= len(eval_loader):
logging.error(f"Batch_metrics ({batch_metrics}) is smaller than batch size "
f"{len(eval_loader)}. Metrics in validation loop won't be computed")
if (batch_index + 1) % batch_metrics == 0: # +1 to skip val loop at very beginning
a, segmentation = torch.max(outputs_flatten, dim=1)
eval_metrics = iou(segmentation, labels_flatten, batch_size, num_classes, eval_metrics)
eval_metrics = report_classification(segmentation, labels_flatten, batch_size, eval_metrics,
ignore_index=eval_loader.dataset.dontcare)
elif dataset == 'tst':
a, segmentation = torch.max(outputs_flatten, dim=1)
eval_metrics = iou(segmentation, labels_flatten, batch_size, num_classes, eval_metrics)
eval_metrics = report_classification(segmentation, labels_flatten, batch_size, eval_metrics,
ignore_index=eval_loader.dataset.dontcare)
logging.debug(OrderedDict(dataset=dataset, loss=f'{eval_metrics["loss"].avg:.4f}'))
if debug and device.type == 'cuda':
res, mem = gpu_stats(device=device.index)
logging.debug(OrderedDict(device=device, gpu_perc=f'{res.gpu} %',
gpu_RAM=f'{mem.used / (1024 ** 2):.0f}/{mem.total / (1024 ** 2):.0f} MiB'))
logging.info(f"{dataset} Loss: {eval_metrics['loss'].avg}")
if batch_metrics is not None:
logging.info(f"{dataset} precision: {eval_metrics['precision'].avg}")
logging.info(f"{dataset} recall: {eval_metrics['recall'].avg}")
logging.info(f"{dataset} fscore: {eval_metrics['fscore'].avg}")
logging.info(f"{dataset} iou: {eval_metrics['iou'].avg}")
return eval_metrics
def train(train_loader,
model,
criterion,
optimizer,
scheduler,
num_classes,
batch_size,
ep_idx,
progress_log,
device,
scale,
vis_params,
debug=False
):
"""
Train the model and return the metrics of the training epoch
:param train_loader: training data loader
:param model: model to train
:param criterion: loss criterion
:param optimizer: optimizer to use
:param scheduler: learning rate scheduler
:param num_classes: number of classes
:param batch_size: number of samples to process simultaneously
:param ep_idx: epoch index (for hypertrainer log)
:param progress_log: progress log file (for hypertrainer log)
:param device: device used by pytorch (cpu ou cuda)
:param scale: Scale to which values in sat img have been redefined. Useful during visualization
:param vis_params: (Dict) Parameters useful during visualization
:param debug: (bool) Debug mode
:return: Updated training loss
"""
model.train()
train_metrics = create_metrics_dict(num_classes)
for batch_index, data in enumerate(tqdm(train_loader, desc=f'Iterating train batches with {device.type}')):
progress_log.open('a', buffering=1).write(tsv_line(ep_idx, 'trn', batch_index, len(train_loader), time.time()))
try: # For HPC when device 0 not available. Error: RuntimeError: CUDA error: invalid device ordinal
inputs = data['sat_img'].to(device)
labels = data['map_img'].to(device)
except RuntimeError:
logging.exception(f'Unable to use device {device}. Trying "cuda:0"')
device = torch.device('cuda')
inputs = data['sat_img'].to(device)
labels = data['map_img'].to(device)
# forward
optimizer.zero_grad()
outputs = model(inputs)
# added for torchvision models that output an OrderedDict with outputs in 'out' key.
# More info: https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101/
if isinstance(outputs, OrderedDict):
outputs = outputs['out']
# vis_batch_range: range of batches to perform visualization on. see README.md for more info.
# vis_at_eval: (bool) if True, will perform visualization at eval time, as long as vis_batch_range is valid
if vis_params['vis_batch_range'] and vis_params['vis_at_train']:
min_vis_batch, max_vis_batch, increment = vis_params['vis_batch_range']
if batch_index in range(min_vis_batch, max_vis_batch, increment):
vis_path = progress_log.parent.joinpath('visualization')
if ep_idx == 0:
logging.info(f'Visualizing on train outputs for batches in range {vis_params["vis_batch_range"]}. '
f'All images will be saved to {vis_path}\n')
vis_from_batch(vis_params, inputs, outputs,
batch_index=batch_index,
vis_path=vis_path,
labels=labels,
dataset='trn',
ep_num=ep_idx + 1,
scale=scale)
loss = criterion(outputs, labels)
train_metrics['loss'].update(loss.item(), batch_size)
if device.type == 'cuda' and debug:
res, mem = gpu_stats(device=device.index)
logging.debug(OrderedDict(trn_loss=f'{train_metrics["loss"].val:.2f}',
gpu_perc=f'{res.gpu} %',
gpu_RAM=f'{mem.used / (1024 ** 2):.0f}/{mem.total / (1024 ** 2):.0f} MiB',
lr=optimizer.param_groups[0]['lr'],
img=data['sat_img'].numpy().shape,
smpl=data['map_img'].numpy().shape,
bs=batch_size,
out_vals=np.unique(outputs[0].argmax(dim=0).detach().cpu().numpy()),
gt_vals=np.unique(labels[0].detach().cpu().numpy())))
loss.backward()
optimizer.step()
scheduler.step()
if train_metrics["loss"].avg is not None:
logging.info(f'Training Loss: {train_metrics["loss"].avg:.4f}')
return train_metrics
def evaluation(eval_loader, model, criterion, num_classes, batch_size, ep_idx, progress_log, batch_metrics=None, dataset='val', device=None, debug=False):
"""
Evaluate the model and return the updated metrics
:param eval_loader: data loader
:param model: model to evaluate
:param criterion: loss criterion
:param num_classes: number of classes
:param batch_size: number of samples to process simultaneously
:param ep_idx: epoch index (for hypertrainer log)
:param progress_log: progress log file (for hypertrainer log)
:param batch_metrics: (int) Metrics computed every (int) batches. If left blank, will not perform metrics.
:param dataset: (str) 'val or 'tst'
:param device: device used by pytorch (cpu ou cuda)
:return: (dict) eval_metrics
"""
eval_metrics = create_metrics_dict(num_classes)
model.eval()
with tqdm(eval_loader, dynamic_ncols=True, desc=f'Iterating {dataset} batches with {device.type}') as _tqdm:
for batch_index, data in enumerate(_tqdm):
progress_log.open('a', buffering=1).write(tsv_line(ep_idx, dataset, batch_index, len(eval_loader), time.time()))
with torch.no_grad():
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
labels_flatten = labels
outputs = model(inputs)
outputs_flatten = outputs
loss = criterion(outputs, labels)
eval_metrics['loss'].update(loss.item(), batch_size)
if (dataset == 'val') and (batch_metrics is not None):
# Compute metrics every n batches. Time consuming.
assert batch_metrics <= len(_tqdm), f"Batch_metrics ({batch_metrics} is smaller than batch size " \
f"{len(_tqdm)}. Metrics in validation loop won't be computed"
if (batch_index+1) % batch_metrics == 0: # +1 to skip val loop at very beginning
a, segmentation = torch.max(outputs_flatten, dim=1)
eval_metrics = report_classification(segmentation, labels_flatten, batch_size, eval_metrics,
ignore_index=get_key_def("ignore_index", params["training"], None))
elif dataset == 'tst':
a, segmentation = torch.max(outputs_flatten, dim=1)
eval_metrics = report_classification(segmentation, labels_flatten, batch_size, eval_metrics,
ignore_index=get_key_def("ignore_index", params["training"], None))
_tqdm.set_postfix(OrderedDict(dataset=dataset, loss=f'{eval_metrics["loss"].avg:.4f}'))
if debug and device.type == 'cuda':
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(OrderedDict(device=device, gpu_perc=f'{res.gpu} %',
gpu_RAM=f'{mem.used/(1024**2):.0f}/{mem.total/(1024**2):.0f} MiB'))
logging.info(f"{dataset} Loss: {eval_metrics['loss'].avg}")
if batch_metrics is not None:
logging.info(f"{dataset} precision: {eval_metrics['precision'].avg}")
logging.info(f"{dataset} recall: {eval_metrics['recall'].avg}")
logging.info(f"{dataset} fscore: {eval_metrics['fscore'].avg}")
return eval_metrics
def sem_seg_inference(model,
nd_array,
overlay,
chunk_size,
num_classes,
device,
meta_map=None,
metadata=None):
"""Inference on images using semantic segmentation
Args:
model: model to use for inference
nd_array: nd_array
overlay: amount of overlay to apply
num_classes: number of different classes that may be predicted by the model
device: device used by pytorch (cpu ou cuda)
returns a numpy array of the same size (h,w) as the input image, where each value is the predicted output.
"""
# switch to evaluate mode
model.eval()
if len(nd_array.shape) == 3:
h, w, nb = nd_array.shape
padded_array = np.pad(nd_array, ((overlay, chunk_size),
(overlay, chunk_size), (0, 0)),
mode='constant')
elif len(nd_array.shape) == 2:
h, w = nd_array.shape
padded_array = np.expand_dims(np.pad(nd_array, ((overlay, chunk_size),
(overlay, chunk_size)),
mode='constant'),
axis=0)
else:
h = 0
w = 0
padded_array = None
output_probs = np.empty(
[num_classes, h + overlay + chunk_size, w + overlay + chunk_size],
dtype=np.float32)
output_counts = np.zeros([output_probs.shape[1], output_probs.shape[2]],
dtype=np.int32)
if padded_array.any():
with torch.no_grad():
with tqdm(range(overlay, h, chunk_size - overlay),
position=1,
leave=False) as _tqdm:
for row in _tqdm:
row_start = row - overlay
row_end = row_start + chunk_size
for col in range(overlay, w, chunk_size - overlay):
col_start = col - overlay
col_end = col_start + chunk_size
chunk_input = padded_array[row_start:row_end,
col_start:col_end, :]
if meta_map:
chunk_input = MetaSegmentationDataset.append_meta_layers(
chunk_input, meta_map, metadata)
inputs = torch.from_numpy(
np.float32(np.transpose(chunk_input, (2, 0, 1))))
inputs.unsqueeze_(0)
inputs = inputs.to(device)
# forward
outputs = model(inputs)
# torchvision models give output it 'out' key. May cause problems in future versions of torchvision.
if isinstance(outputs,
OrderedDict) and 'out' in outputs.keys():
outputs = outputs['out']
output_counts[row_start:row_end,
col_start:col_end] += 1
output_probs[:, row_start:row_end,
col_start:col_end] += np.squeeze(
outputs.cpu().numpy(), axis=0)
if debug and device.type == 'cuda':
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(
OrderedDict(
device=device,
gpu_perc=f'{res.gpu} %',
gpu_RAM=
f'{mem.used / (1024 ** 2):.0f}/{mem.total / (1024 ** 2):.0f} MiB',
chunk_size=inputs.cpu().numpy().shape,
output_size=outputs.cpu().numpy().shape))
output_mask = np.argmax(np.divide(output_probs,
np.maximum(output_counts, 1)),
axis=0)
# Resize the output array to the size of the input image and write it
return output_mask[overlay:(h + overlay),
overlay:(w + overlay)].astype(np.uint8)
else:
raise IOError(
f"Error classifying image : Image shape of {len(nd_array.shape)} is not recognized"
)
def evaluation(eval_loader,
model,
criterion,
num_classes,
batch_size,
task,
ep_idx,
progress_log,
batch_metrics=None,
dataset='val',
device=None):
"""
Evaluate the model and return the updated metrics
:param eval_loader: data loader
:param model: model to evaluate
:param criterion: loss criterion
:param num_classes: number of classes
:param batch_size: number of samples to process simultaneously
:param task: segmentation or classification
:param ep_idx: epoch index (for hypertrainer log)
:param progress_log: progress log file (for hypertrainer log)
:param batch_metrics: (int) Metrics computed every (int) batches. If left blank, will not perform metrics.
:param dataset: (str) 'val or 'tst'
:param device: device used by pytorch (cpu ou cuda)
:return: (dict) eval_metrics
"""
eval_metrics = create_metrics_dict(num_classes)
model.eval()
with tqdm(eval_loader, dynamic_ncols=True) as _tqdm:
for index, data in enumerate(_tqdm):
progress_log.open('a', buffering=1).write(
tsv_line(ep_idx, dataset, index, len(eval_loader),
time.time()))
with torch.no_grad():
if task == 'classification':
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
labels_flatten = labels
outputs = model(inputs)
outputs_flatten = outputs
elif task == 'segmentation':
inputs = data['sat_img'].to(device)
labels = data['map_img'].to(device)
labels_flatten = flatten_labels(labels)
outputs = model(inputs)
if isinstance(outputs, OrderedDict):
outputs = outputs['out']
outputs_flatten = flatten_outputs(outputs, num_classes)
loss = criterion(outputs, labels)
eval_metrics['loss'].update(loss.item(), batch_size)
if (dataset == 'val') and (batch_metrics is not None):
# Compute metrics every n batches. Time consuming.
if index + 1 % batch_metrics == 0: # +1 to skip val loop at very beginning
a, segmentation = torch.max(outputs_flatten, dim=1)
eval_metrics = report_classification(
segmentation, labels_flatten, batch_size,
eval_metrics)
elif dataset == 'tst':
a, segmentation = torch.max(outputs_flatten, dim=1)
eval_metrics = report_classification(
segmentation, labels_flatten, batch_size, eval_metrics)
_tqdm.set_postfix(
OrderedDict(dataset=dataset,
loss=f'{eval_metrics["loss"].avg:.4f}'))
if debug and torch.cuda.is_available(
) and device.type != "cpu":
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(
OrderedDict(
device=device,
gpu_perc=f'{res.gpu} %',
gpu_RAM=
f'{mem.used/(1024**2):.0f}/{mem.total/(1024**2):.0f} MiB'
))
print(f"{dataset} Loss: {eval_metrics['loss'].avg}")
if batch_metrics is not None:
print(f"{dataset} precision: {eval_metrics['precision'].avg}")
print(f"{dataset} recall: {eval_metrics['recall'].avg}")
print(f"{dataset} fscore: {eval_metrics['fscore'].avg}")
return eval_metrics
def train(train_loader, model, criterion, optimizer, scheduler, num_classes,
batch_size, task, ep_idx, progress_log, device):
"""
Train the model and return the metrics of the training epoch
:param train_loader: training data loader
:param model: model to train
:param criterion: loss criterion
:param optimizer: optimizer to use
:param scheduler: learning rate scheduler
:param num_classes: number of classes
:param batch_size: number of samples to process simultaneously
:param task: segmentation or classification
:param ep_idx: epoch index (for hypertrainer log)
:param progress_log: progress log file (for hypertrainer log)
:param device: device used by pytorch (cpu ou cuda)
:return: Updated training loss
"""
model.train()
train_metrics = create_metrics_dict(num_classes)
with tqdm(train_loader) as _tqdm:
for index, data in enumerate(_tqdm):
progress_log.open('a', buffering=1).write(
tsv_line(ep_idx, 'trn', index, len(train_loader), time.time()))
if task == 'classification':
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
elif task == 'segmentation':
inputs = data['sat_img'].to(device)
labels = data['map_img'].to(device)
# forward
optimizer.zero_grad()
outputs = model(inputs)
# added for torchvision models that output an OrderedDict with outputs in 'out' key.
# More info: https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101/
if isinstance(outputs, OrderedDict):
outputs = outputs['out']
loss = criterion(outputs, labels)
train_metrics['loss'].update(loss.item(), batch_size)
if debug and device.type == 'cuda':
res, mem = gpu_stats(device=device.index)
_tqdm.set_postfix(
OrderedDict(
train_loss=f'{train_metrics["loss"].val:.4f}',
device=device,
gpu_perc=f'{res.gpu} %',
gpu_RAM=
f'{mem.used/(1024**2):.0f}/{mem.total/(1024**2):.0f} MiB',
learning_rate=optimizer.param_groups[0]['lr'],
img_size=data['sat_img'].numpy().shape,
sample_size=data['map_img'].numpy().shape,
batch_size=batch_size))
loss.backward()
optimizer.step()
scheduler.step()
print(f'Training Loss: {train_metrics["loss"].avg:.4f}')
return train_metrics
