Реализация задачи сегментации ячеек ISBI на основе U-образной сети с OneFlow

алгоритм
目录
1.简介
2.网路架构
3.数据和程序准备
4.使用步骤
5.单机单卡训练方式
6.单机多卡训练方式(DDP)
7.可视化实验结果
8.小结

Автор Ли Сян

1

Введение

Для пользователей, знакомых с PyTorch (таких как я), вы можете быстро начать работу с OneFlow, что очень удобно. В этом документе реализована сегментация ячеек ISBI-вызова на основе OneFlow и U-Net.Код включает два метода обучения: одна машина с одной картой и одна машина с несколькими картами. Среди них одномашинное многокарточное обучение — это параллелизм данных с помощью модуля oneflow.nn.parallel.DistributedDataParallel и лаунчера.

Кроме того, поскольку в настоящее время я являюсь стажером по алгоритмам в OneFlow, в этой статье будет проанализировано больше с точки зрения пользователя, который плохо знаком с инфраструктурой OneFlow, включая API, возможности распределенного обучения, высокую производительность и часть моего опыта стажировки.

Задача сегментации клеток ISBI: учитывая карту структуры клеток, контуры краев классифицируются по двум категориям, как показано на следующей анимации.

Имеется 30 обучающих данных с разрешением 512x512.Эти изображения представляют собой электронно-микроскопические карты клеток дрозофилы.

2

сетевая архитектура

Архитектура сети U-Net показана на следующем рисунке. Он состоит из сокращающегося пути и расширяющегося пути, а сокращающийся путь соответствует типичной архитектуре сверточных сетей. Он состоит из повторного использования двух сверток 3x3, за каждой из которых следует блок линейной коррекции (ReLU), и операции максимального объединения 2x2 с понижающей дискретизацией с шагом 2.

На каждом этапе понижения дискретизации мы удваиваем количество функциональных каналов. Каждый шаг на пути расширения состоит из повышающей дискретизации карты объектов, за которой следует повышающая свертка 2x2, уменьшающая вдвое количество каналов объектов, в сочетании с соответствующей обрезанной картой объектов из пути сокращения. Затем две свертки 3x3, за каждой следует ReLU. Обрезка необходима, потому что граничные пиксели теряются при каждой свертке. На последнем слое вектор признаков каждого компонента сопоставляется с желаемым количеством классов (2 класса) с использованием сверток 1x1. Всего в сети 23 сверточных слоя.

В соответствии с приведенной выше структурой сети использование OneFlow для реализации кода U-образной структуры сети выглядит следующим образом:

"""
Creates a U-Net Model as defined in:
U-Net: Convolutional Networks for Biomedical Image Segmentation
https://arxiv.org/abs/1505.04597
Modified from https://github.com/milesial/Pytorch-UNet
"""
import oneflow as flow
import oneflow.nn as nn
import oneflow.nn.functional as F


class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2), DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(
                in_channels, out_channels, kernel_size=2, stride=2
            )

        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2))

        x = flow.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)


class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 1024)
        self.up1 = Up(1024, 512, bilinear)
        self.up2 = Up(512, 256, bilinear)
        self.up3 = Up(256, 128, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

3

Подготовка данных и программы

Необработанные данные: сначала подготовьте данные, справочные данные взяты из набора данных ISBI. Данные могут храниться в этом репозитории (*GitHub.com/oneflow-Inc…

Расширенные данные: Google Cloud Drive (drive.Google.com/drive/Фландрия…)

Выберите один из вышеуказанных данных.

Ссылка на код:GitHub.com/oneflow-Inc…

Каталог программы выглядит следующим образом:

dataloader.py//加载数据
plot.py//绘制loss曲线
TrainUnetDataSet.py//训练文件
unet.py//网路结构
predict_unet_test.py//测试文件
tran.sh//训练脚本
test.sh//测试脚本

4

Шаги для использования

тренироваться:

bash train.sh

контрольная работа:

bash test.sh

5

Метод обучения с одной машиной и одной картой

В TrainUnetDataSet.py, чтобы сравнить с методом обучения с несколькими картами на одной машине, вот полный сценарий для обучения U-Net, как показано ниже:

def Train_Unet(net, device, data_path, batch_size=3, epochs=40, lr=0.0001):
    train_dataset = SelfDataSet(data_path)
    train_loader = utils.data.DataLoader(
        train_dataset, batch_size=batch_size, shuffle=True
    )

    opt = optim.Adam((net.parameters()))
    loss_fun = nn.BCEWithLogitsLoss()
    bes_los = float("inf")

    for epoch in range(epochs):
        net.train()
        running_loss = 0.0
        i = 0
        begin = time.perf_counter()
        for image, label in train_loader:
            opt.zero_grad()
            image = image.to(device=device, dtype=flow.float32)
            label = label.to(device=device, dtype=flow.float32)
            pred = net(image)
            loss = loss_fun(pred, label)
            loss.backward()
            i = i + 1
            running_loss = running_loss + loss.item()
            opt.step()
        end = time.perf_counter()
        loss_avg_epoch = running_loss / i
        Unet_train_txt.write(str(format(loss_avg_epoch, ".4f")) + "\n")
        print("epoch: %d avg loss: %f time:%d s" % (epoch, loss_avg_epoch, end - begin))
        if loss_avg_epoch < bes_los:
            bes_los = loss_avg_epoch
            state = {"net": net.state_dict(), "opt": opt.state_dict(), "epoch": epoch}
            flow.save(state, "./checkpoints")

def main(args):
    DEVICE = "cuda" if flow.cuda.is_available() else "cpu"
    print("Using {} device".format(DEVICE))
    net = UNet(1, 1, bilinear=False)
    # print(net)
    net.to(device=DEVICE)
    data_path = args.data_path
    Train_Unet(net, DEVICE, data_path, epochs=args.epochs, batch_size=args.batch_size)
    Unet_train_txt.close()

6

Метод многокарточного обучения на одной машине (DDP)

Перспектива SBP и согласованности, предоставляемая OneFlow, упрощает сложные конфигурации распределенного обучения, но я все еще учусь и практикуюсь (научитесь делиться в следующий раз). Поскольку здесь требуется только простое параллельное обучение данных, я использовал модуль oneflow.nn.parallel.DistributedDataParallel и средство запуска, совмещенное с torch.nn.parallel.DistributedDataParallel, чтобы оно выполнялось без каких-либо изменений в одномашинном однопроцессорном режиме. карточный скрипт Параллельное обучение данных.

В соответствии с этой функцией существует только два различия между кодом обучения с параллельными данными и сценарием с одной картой для одной машины.Измените сценарий обучения в разделе 5 следующим образом:

1. Используйте DistributedDataParallel для обработки объекта модуля.

    m=net.to(device=DEVICE)
    net = ddp(m)

2. Используйте DistributedSampler для создания экземпляра загрузчика данных в каждом процессе, каждый экземпляр загрузчика данных загружает часть полных данных и автоматически завершает распределение данных.

    is_distributed=True
    sampler = flow.utils.data.distributed.DistributedSampler(train_dataset) if is_distributed else None
    train_loader = utils.data.DataLoader(
        train_dataset, batch_size=batch_size, shuffle=(sampler is None), sampler=sampler
    )

В распределенном режиме метод set_epoch() вызывается в начале каждой эпохи перед созданием итератора DataLoader, необходимого для правильной перетасовки между несколькими эпохами. В противном случае всегда будет использоваться один и тот же порядок.

    for epoch in range(epochs):
        if is_distributed:
            sampler.set_epoch(epoch)
            ···

На этом написание сценария распределенного обучения завершено, затем используйте средство запуска для запуска сценария, а остальное оставьте OneFlow, чтобы распределенное обучение U-Net было таким же простым, как обучение U-Net с одной машиной и одним карта.

python3 -m oneflow.distributed.launch --nproc_per_node 8 ./TrainUnetDataSet.py

Параметр --nproc_per_node указывает количество вызываемых узлов GPU.

7

Визуализируйте экспериментальные результаты

В этом эксперименте обучается только 40 эпох, и увеличение количества эпох может значительно повысить точность модели.

Исходное изображение:

Исходное изображение для сегментации

График предсказания U-Net:

отображение диаграммы предсказания unet

8

резюме

Эта статья больше проанализирована с точки зрения пользователя, который впервые столкнулся с инфраструктурой OneFlow.В процессе использования OneFlow для обучения сети U-Net я действительно почувствовал, что OneFlow дружелюбен к новым пользователям.

В предыдущем исследовании я в основном использовал Keras и TensorFlow, но я могу быстро приступить к работе с OneFlow. Из-за режима Eager OneFlow он согласован с PyTorch, что позволяет пользователям, знакомым с PyTorch, начать работу напрямую без каких-либо затрат. Что касается режима Graph, я еще не практиковал его, но Graph также основан на объектно-ориентированном стиле программирования.Пользователи, знакомые с разработкой динамических графов, могут использовать эффективные статические графы с небольшими изменениями кода.

Прежде всего, API, предоставляемый OneFlow, в основном может удовлетворить все мои потребности. В следующей версии будут предоставлены более подробные справочные документы API и более богатый и полный набор операторов. По сравнению со сложными и большими документами TensorFlow, я думаю, что OneFlow более мощный Простота использования.

Кроме того, производительность OneFlow на крупномасштабных моделях является наиболее важной. Масштабы моделей в области компьютерного зрения становятся все больше и больше, а метод многоузловых кластеров для распределенного обучения для повышения вычислительной мощности лучше решает OneFlow. Более того, простая работа с распределенным обучением тоже больше подходит для моих нужд.

Напоследок расскажу о короткой стажировке в OneFlow: я присоединился к ней 15 ноября, с тех пор прошел месяц. В дополнение к ощущению того, что внутренняя структура глубокого обучения появляется благодаря словам и делам брата BBuf Xiaoyu, я также чувствую эффективный метод разработки команды OneFlow. Месяц пролетел быстро.По сравнению с моим собственным отсутствием пайплайнового обучения, учеба и работа в OneFlow действительно очень полезны для личного совершенствования.Каждый день стажировки очень приятен.

Добро пожаловать, чтобы загрузить и испытать новое поколение среды глубокого обучения OneFlow с открытым исходным кодом:GitHub.com/oneflow-Inc…