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

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//測試指令碼

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

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)

···