網(wǎng)絡(luò)結(jié)構(gòu)

U-net是MICCAI2015的專門針對醫(yī)學(xué)影像分割設(shè)計的網(wǎng)絡(luò)結(jié)構(gòu)，直到今天還是很好用，效果任然很好。網(wǎng)絡(luò)的左半部分是Encode
部分，此部分將輸入圖像逐步下采樣得到高層的語義信息。右半部分是Decoder部分，將Encode提取的語義特征逐步恢復(fù)為原始圖像的尺寸，中間是跳躍連接，將下采樣過程中提取到的不同level的特征加入到Decoder中，得到更好的分割結(jié)果。

完整訓(xùn)練代碼和數(shù)據(jù)可以去我的github上面下載：https://github.com/SPECTRELWF/Semantic-Segmentation
個人主頁：liuweifeng.top:8090

數(shù)據(jù)集介紹

數(shù)據(jù)集是來自網(wǎng)上的視網(wǎng)膜血管分割數(shù)據(jù)集，訓(xùn)練集只有二十張圖像，少得可憐。測試集也是二十張圖像。原始數(shù)據(jù)集中image和label使用.tif和.gif格式存儲的，需要先進(jìn)行一下預(yù)處理，將數(shù)據(jù)組織成想要的樣子，再來進(jìn)行訓(xùn)練。

train

# !/usr/bin/python3 # -*- coding:utf-8 -*- # Author:WeiFeng Liu # @Time: 2021/12/9 下午1:54from unet_model.unet import UNET from dataset import SWXG_Dataset import torch.optim as optim import torch.nn as nn import torchdef train_net(net,device,data_path,epochs=40,batch_size=1,lr=1e-5):isbi_dataset = SWXG_Dataset(data_path)train_loader = torch.utils.data.DataLoader(isbi_dataset,batch_size,shuffle = True)#使用RMSprop優(yōu)化optimizer = optim.RMSprop(net.parameters(),lr,weight_decay=1e-8,momentum=0.9)criterion = nn.BCEWithLogitsLoss()best_loss = float("inf")for epoch in range(epochs):net.train()for images, labels in train_loader:optimizer.zero_grad()images = images.to(device,dtype = torch.float32)labels = labels.to(device,dtype=torch.float32)pred = net(images)loss = criterion(pred,labels)print('epoch:%d train loss:%f' % (epoch+1,loss.item()))if loss <best_loss:best_loss = losstorch.save(net.state_dict(), 'best_model.pth')loss.backward()optimizer.step()if __name__ == "__main__":device = torch.device('cuda:0')net = UNET(n_channels=1,n_classes=1)net.to(device)data_path = r'Dataset/train/'train_net(net,device,data_path)

結(jié)果：

訓(xùn)練30個epoch之后的結(jié)果看起來還是很好的。

總結(jié)

以上是生活随笔為你收集整理的Pytorch实现U-net视网膜血管分割的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。