用OneFlow實現基於U型網路的ISBI細胞分割任務

語言: CN / TW / HK
目錄 1.簡介 2.網路架構 3.資料和程式準備 4.使用步驟 5.單機單卡訓練方式 6.單機多卡訓練方式(DDP) 7.視覺化實驗結果 8.小結

撰文 | 李響

1

簡介

對於熟悉 PyTorch 的使用者(比如我),可以快速上手 OneFlow,非常絲滑。 本文基於OneFlow和U-Net實現ISBI挑戰賽的細胞分割,程式碼包括單機單卡和單機多卡兩種訓練方式。其中,單機多卡的訓練 是藉助   oneflow.nn.parallel.DistributedDataParallel   模組及   launcher 做的資料並行。

此外,由於我目前在OneFlow做一名演算法實習生,本文更多以一個初次接觸OneFlow框架的使用者角度進行分析,包括API、分散式訓練能力、高效能和我的一些實習感受(:eyes:)。

ISBI細胞分割任務:給一張細胞結構圖,對邊緣輪廓進行二分類,如下動圖所示。

訓練資料有30張,解析度為 512x512 ,這些圖片是果蠅的電鏡細胞圖。

2

網路架構

U-Net網路架構如下圖所示。它由一個收縮路徑和一個擴充套件路徑組成,收縮路徑遵循卷積網路的典型架構。它包括重複使用兩個 3x3 卷積,每個卷積後跟一個線性修正單元(ReLU)和一個2x2最大池化操作,步長為2的下采樣。

在每個下采樣步驟中,我們將特徵通道的數量加倍。擴充套件路徑中的每一步都包括特徵對映的上取樣,然後進行 2x2 向上卷積,將特徵通道數量減半,與來自收縮路徑的相應裁剪特徵對映串聯。然後是兩個3x3卷積,每個卷積後面接ReLU。由於每一次卷積都會丟失邊界畫素,因此裁剪是必要的。在最後一層,使用1x1卷積將每個分量特徵向量對映到所需數量的類別(2類)上。網路總共有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( 64128 )

self.down2 = Down( 128256 )

self.down3 = Down( 256512 )

self.down4 = Down( 5121024 )

self.up1 = Up( 1024512 , bilinear)

self.up2 = Up( 512256 , bilinear)

self.up3 = Up( 256128 , bilinear)

self.up4 = Up( 12864 , 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 挑戰的資料集。資料可以在本倉庫 ( https://github.com/Oneflow-Inc/models/tree/main/Vision/segmentation/U-Net )下載到,含30張訓練圖、30張對應的標籤。30張測試圖片。

增強後的資料 :谷歌雲盤( https://drive.google.com/drive/folders/0BzWLLyI3R0pLclMzMzRxUm1qZmc )

以上資料二選一。

程式碼連結: https://github.com/Oneflow-Inc/models/tree/main/Vision/segmentation/U-Net

該程式目錄如下:

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

OneFlow 提供的 SBP 和一致性視角,讓複雜的分散式訓練配置變得更簡單,不過我還在學習實踐中(學會了下次分享)。因為這裡的需求只是簡單的資料並行訓練,所以我使用了與 torch.nn.parallel.DistributedDataParallel   對齊的  oneflow.nn.parallel.DistributedDataParallel  模組及 launcher ,這樣幾乎不用對單機單卡指令碼做修改,就完成了資料並行訓練。

根據該特性,資料並行的訓練程式碼與單機單卡指令碼的不同只有2個,將第5節的訓練指令碼做如下修改:

1.使用  DistributedDataParallel  處理一下 module 物件

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

2.使用 DistributedSampler 在每個程序中例項化 Dataloader ,每個 Dataloader 例項載入完整資料的一部分,自動完成資料的分發。

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

)

分散式模式下,在建立DataLoader迭代器之前,在每個epoch開始時呼叫set_epoch()方法,這對在多個epoch中正確地進行shuffle是必要的。否則,將總是使用相同的順序。

for epoch  in range(epochs):

if is_distributed:

sampler.set_epoch(epoch)

···

這樣就完成了分散式訓練指令碼的編寫,然後使用  launcher  啟動指令碼,把剩下的一切都交給 OneFlow,讓分散式訓練U-Net,像單機單卡訓練U-Net一樣簡單。

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

--nproc_per_node 選項表示呼叫的GPU結點數量。

7

視覺化實驗結果

該實驗只訓練了40個Epoch,提升Epoch數量可顯著提升模型精度。

原圖:

原始待分割圖片

U-Net預測圖:

unet預測圖展示

8

小結

本文更多以一個初次接觸OneFlow框架的使用者角度進行分析,在使用OneFlow訓練U-Net網路的過程中,我真正感受到OneFlow對新使用者的友好。

之前的學習中,我主要使用Keras和TensorFlow,但使用OneFlow卻可以很快上手。由於OneFlow的Eager模式,與PyTorch對齊,讓熟悉PyTorch的使用者可以零成本直接上手。至於Graph模式,目前我還沒有進行實踐,但Graph也是基於面向物件的程式設計風格,熟悉動態圖開發的使用者,只需要改很少量的程式碼,就可以使用高效率的靜態圖。

首先,OneFlow提供的API基本可以滿足我的所有需求,在下一版本中也將提供更加詳細的API幫助文件和更豐富、完善的運算元集,對比TensorFlow複雜和大型的文件,我認為OneFlow更具有易用性。

此外,OneFlow在處理大規模模型上的效能是最重要的。計算機視覺領域的模型規模越來越大,多節點叢集進行分散式訓練,以提升算力的方法被OneFlow更好地解決了。而且,分散式訓練的簡單操作也更能滿足我的需求。

最後,說一下在OneFlow短暫的實習體驗:11月15日入職,到現在為止已經過去一個月了。除了感受到國產深度學習框架正在異軍突起之外,在BBuf曉雨哥的言傳身教下,也感受到OneFlow團隊的高效開發方式。一個月的時間過得很快,對比我自己沒有pipline式的學習,在OneFlow學習和工作確實對個人提升很有benefits,實習的每天過得很愉快。

其他人都在看

點選“ 閱讀原文 ,歡迎下載體驗OneFlow新一代開源深度學習框架