最近由於項目需要做了一段時間的語義分割，希望能將自己的心路歷程記錄下來，以提供給所需幫助的人
接下來我將依託Unet語義分割網絡介紹以下內容：
首先我的環境配置
pytorch1.10
win10
vs2017
python3.6
opencv3.4
Aaconda-5.2.0
一、使用pytorch實現簡單的unet分割網絡
二、使用Unet做多類別分割
三、c++調用python執行語義分割
四、c++調用libtorch執行語義分割
首先介紹使用pytorch實現簡單的unet分割網絡
Unet在醫學圖像分割領域是一個比較有名的分割網絡，用Unet入坑語義分割由以下幾個優點：
1、由於使用Unet模型做語義分割不需要大量的標註樣本，甚至幾十張標註樣本都能取得較好的效果。
2、短時間內就能獲取訓練效果，可以根據自己的想法隨意DIY Unet網絡。
用pytorch實現Unet網絡也很簡單，下面是我復現的一個簡單的Unet分割模型：
class conv_bn_relu(nn.Module):
def __init__(self, in_ch, out_ch,kerl = 3,pad = 1 ):
super(conv_bn_relu, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_ch, out_ch, kernel_size = kerl, padding = pad),
nn.BatchNorm2d(out_ch, eps=1e-3),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, kernel_size = kerl, padding = pad),
nn.BatchNorm2d(out_ch, eps=1e-3),
nn.ReLU(inplace=True)
)
def forward(self, input):
c = self.conv(input)
return c
class Unet(nn.Module):
def __init__(self,in_ch,out_ch):
super(Unet, self).__init__()
self.conv1 = conv_bn_relu(in_ch, 64)
self.pool1 = nn.MaxPool2d(2)
self.conv2 = conv_bn_relu(64, 128)
self.pool2 = nn.MaxPool2d(2)
self.conv3 = conv_bn_relu(128, 256)
self.up4 = nn.ConvTranspose2d(256, 128, 2, stride=2)
self.conv4 = conv_bn_relu(256, 128)
self.up5 = nn.ConvTranspose2d(128, 64, 2, stride=2)
self.conv5 = conv_bn_relu(128, 64)
self.conv6 = nn.Conv2d(64,out_ch, 1)
def forward(self,x):
c1=self.conv1(x)
p1=self.pool1(c1)
c2=self.conv2(p1)
p2=self.pool2(c2)
c3=self.conv3(p2)
up_4=self.up5(c3)
merge4 = torch.cat([up_4, c2], dim=1)
c4=self.conv4(merge4)
up_5=self.up4(c4)
merge5=torch.cat([up_5, c1],dim=1)
c5=self.conv5(merge5)
c6=self.conv6(c5)
return c6
顯然Unet分割網絡比其他語義分割模型更簡潔，降低了小白入坑的門檻。
有了模型後我們如何去訓練一個自己的模型呢，首先也是重要的一點就是數據讀取，關於讀取數據我寫了一個簡單的例子：
首先是獲取樣本圖像的原始圖像以及標註圖像的路徑
def make_dataset(root):
imgs=[]
path = os.path.join(root, "001.txt")
file = open(path)
for line in file.readlines():
line = line.split('\n')[0] #特別注意，去除文件路徑後面的空白部分
img = os.path.join(root, "data", line)
mask = os.path.join(root, "label", line)
imgs.append((img,mask))
return imgs
我把原始數據與標註數據放在不同的文件夾下，且原始數據與標註數據的名稱都相同。
所以文windows系統下只需要批處理生成001.txt文件就行了，具體操作如下：
首先在原始圖像文件夾下建立一個test.txt文件，裏面寫上：
dir /b *.png *.jpg *.raw *.tif > 001.txt 這是相對路徑
生成的只有圖片文件名，形如 1.png
dir /s /b *.png *.jpg *.raw *.tif > 001.txt 這是絕對路徑
生成的絕對路徑，形如 D:\data\1.png
我們的txt文件中只需要寫這一句 dir /b *.png *.jpg *.raw *.tif > 001.txt ，然後將test.txt文件名的後綴改爲test.bat,然後雙擊運行就可以獲得我們所需的txt文件
然後就是數據加載部分了，數據加載繼承於Dataset，一般寫法如下：
class LoadDataset(Dataset):
def __init__(self, root, transform=None, target_transform=None):
imgs = make_dataset(root)
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
def __getitem__(self, index):
x_path, y_path = self.imgs[index]
img_x = cv2.imread(x_path,0)
img_y = cv2.imread(y_path,0)
if self.transform is not None:
img_x = self.transform(img_x)
if self.target_transform is not None:
img_y = self.target_transform(img_y)
return img_x, img_y
def __len__(self):
return len(self.imgs)
接下來就輪到訓練部分了
首先寫上這兩行，使得訓練更穩定，具體原因可百度
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
原始數據與標註數據準換成Tensor
x_transforms = transforms.Compose([
transforms.ToTensor(),
#transforms.Normalize(mean=[0.85], std=[0.12]) #可根據需要決定
])
y_transforms = transforms.ToTensor()
訓練模塊如下：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def train_model(model, criterion, optimizer, dataload, num_epochs = 100):
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
dt_size = len(dataload.dataset)
epoch_loss = 0
step = 0
for x, y in dataload:
step += 1
inputs = x.to(device)
labels = y.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
epoch_loss += float(loss.item())
print("%d/%d,train_loss:%0.5f" % (step, (dt_size - 1) // dataload.batch_size + 1, loss.item()))
torch.cuda.empty_cache() #及時清除無用緩存
print("epoch %d loss:%0.5f batch size " % (epoch, epoch_loss/step))
if np.isnan(epoch_loss/step):
break
torch.save(model.state_dict(), str(epoch)+ "_" + str(epoch_loss/step) + '_temp.pth') #保存每一輪的結果
torch.save(model.state_dict(), 'Unet.pth')
return model
def train():
model = Unet(1,1).to(device)
criterion = CrossEntropyLoss2d()
optimizer = optim.Adam(model.parameters(),lr=0.001)
liver_dataset = LoadDataset("D:/u-net/sample/train",transform = x_transforms,target_transform = y_transforms)
dataloaders = DataLoader(liver_dataset, batch_size = 4, shuffle=True, num_workers=4,pin_memory=True)
train_model(model, criterion, optimizer, dataloaders)
模型參數較大的情況下不推薦每輪保存一次模型，優化器只推薦Adam，batch_size視GPU顯存的大小而定，pin_memory=True也需要根據GPU判定是否要開啓（pytorch 默認爲False），關於DataLoader參數設置可自行百度。
損失函數我們有更好的選擇例如FocalLoss：
class FocalLoss(torch.nn.Module):
def __init__(self, gamma=2, alpha=0.25, reduction='elementwise_mean'):
super().__init__()
self.gamma = gamma
self.alpha = alpha
self.reduction = reduction
def forward(self, _input, target):
pt = torch.sigmoid(_input)
alpha = self.alpha
loss = - alpha * (1 - pt) ** self.gamma * target * torch.log(pt) - \
(1 - alpha) * pt ** self.gamma * (1 - target) * torch.log(1 - pt)
if self.reduction == 'elementwise_mean':
loss = torch.mean(loss)
elif self.reduction == 'sum':
loss = torch.sum(loss)
return loss
經測試二值語義分割FocalLoss要優於CrossEntropyLoss2d ，其他loss可自行嘗試。
爲了保證結果的可復現，我們每次訓練手動初始化隨機參數：
def seed_torch(seed=666):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
訓練之前要測試以下數據處理是否正確：
def data_test():
liver_dataset = LoadDataset("D:/u-net/sample/train",transform=x_transforms,target_transform = y_transforms)
dataloaders = DataLoader(liver_dataset, batch_size=1, shuffle=True, num_workers=0)
for x, y in dataloaders:
img_x = np.uint8(torch.squeeze(x).numpy()*255)
img_y = np.uint8(torch.squeeze(y).numpy()*255)
print(img_x.shape)
print(img_y.shape)
cv2.imshow("img_x",img_x)
cv2.imshow("img_y",img_y)
cv2.waitKey(0)
break
另外需要注意：輸入圖像需要爲2的倍數，或者保證所有層池化後反捲積能恢復同樣的大小。
訓練所需數據 可在kaggle下載
下一步篇將介紹如何用Unet做多類別分割