def train(loader, num_classes, device, net, optimizer, criterion):
num_samples = 0
running_loss = 0
metrics = Metrics()
net.train()
for images, masks, _tile in tqdm(loader,
desc="Train",
unit="batch",
ascii=True):
images = images.to(device)
masks = masks.to(device)
assert images.size()[2:] == masks.size(
)[1:], "resolutions for images and masks are in sync"
num_samples += int(images.size(0))
optimizer.zero_grad()
outputs = net(images)
assert outputs.size()[2:] == masks.size(
)[1:], "resolutions for predictions and masks are in sync"
assert outputs.size(
)[1] == num_classes, "classes for predictions and dataset are in sync"
loss = criterion(outputs, masks)
loss.backward()
optimizer.step()
running_loss += loss.item()
for mask, output in zip(masks, outputs):
prediction = output.detach()
metrics.add(mask, prediction)
assert num_samples > 0, "dataset contains training images and labels"
class_stats = metrics.get_classification_stats()
return {
"loss": running_loss / num_samples,
"miou": metrics.get_miou(),
"fg_iou": metrics.get_fg_iou(),
"mcc": metrics.get_mcc(),
"accuracy": class_stats['accuracy'],
"precision": class_stats['precision'],
"recall": class_stats['recall'],
"f1": class_stats['f1']
}
def validate(loader, num_classes, device, net, criterion):
num_samples = 0
running_loss = 0
metrics = Metrics()
net.eval()
with torch.no_grad():
for images, masks, _tile in tqdm(loader,
desc="Validate",
unit="batch",
ascii=True):
images = images.to(device)
masks = masks.to(device)
assert images.size()[2:] == masks.size(
)[1:], "resolutions for images and masks are in sync"
num_samples += int(images.size(0))
outputs = net(images)
assert outputs.size()[2:] == masks.size(
)[1:], "resolutions for predictions and masks are in sync"
assert outputs.size(
)[1] == num_classes, "classes for predictions and dataset are in sync"
loss = criterion(outputs, masks)
running_loss += loss.item()
for mask, output in zip(masks, outputs):
metrics.add(mask, output)
assert num_samples > 0, "dataset contains validation images and labels"
class_stats = metrics.get_classification_stats()
return {
"loss": running_loss / num_samples,
"miou": metrics.get_miou(),
"fg_iou": metrics.get_fg_iou(),
"mcc": metrics.get_mcc(),
"accuracy": class_stats['accuracy'],
"precision": class_stats['precision'],
"recall": class_stats['recall'],
"f1": class_stats['f1']
}