文章目录
1. Introduction
2. 网路架构
3. 数据和程序准备
4. 使用步骤
5. 单机单卡训练方式
6. 单机多卡训练方式(DDP)
7. 可视化实验结果
8. Conclusion and discussion

"""
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

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

bash train.sh

bash test.sh

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()

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

    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
    )

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

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