Пример #1
0
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']
    }
Пример #2
0
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']
    }