移动端unet人像分割模型--1
个人对移动端神经网络开发一直饶有兴致。去年腾讯开源了NCNN框架之后，一直都在关注。近期成功利用别人训练好的mtcnn和mobilefacenet模型制作了一个ios版本人脸识别swift版本demo。希望maskrcnn移植到ncnn，在手机端实现一些有趣的应用。因为unet模型比较简单，干脆就从这个入手。
基本的网络基于keras版本: https://github.com/TianzhongSong/Person-Segmentation-Keras
不过keras没办法直接转成ncnn模型，研究过通过onnx模型做中间跳板，采用了一些开源的转换工具，也是一堆问题。NCNN支持几个神经网络训练框架:caffe/mxnet/pytorch，在ncnn的github有一篇issue里nihui推荐采用mxnet，因此mxnet也成为了我的首选。
利用Person-Segmentation-Keras项目的数据集，同时基于https://github.com/milesial/Pytorch-UNet/tree/master/unet这个项目捣鼓了几段代码。训练完成，用来测试ncnn转换基本可用。
转换过程发现许多问题，一个是调用ncnn extract会crash，经过调查，发现mxnet2ncnn工具也有bug，blob个数算错，其次是input层one_blob_only标志我的理解应该是false，不知道什么原因转换过来的模型这边是true，导致forward_layer函数里面bottoms变量访问异常。后来一层层extract出来打印输出的channel/width/height调查后又发现，我把unet.py里的name为pool5写成了pool4（文章中的code已经纠正），可能前面的crash跟这个致命错误有关系也说不定。只好重新训练模型，几个小时漫长等待，剩下部分下周再写。部分代码已经更新，请参考: https://github.com/xuduo35/unet_mxnet2ncnn
unetdataiter.py
#!/usr/bin/env python
# coding=utf8
import os
import sys
import random
import cv2
import mxnet as mx
import numpy as np
from mxnet.io import DataIter, DataBatch
sys.path.append('../')
def get_batch(items, root_path, nClasses, height, width):
x = []
y = []
for item in items:
image_path = root_path + item.split(' ')[0]
label_path = root_path + item.split(' ')[-1].strip()
img = cv2.imread(image_path, 1)
label_img = cv2.imread(label_path, 1)
im = np.zeros((width, height, 3), dtype='uint8')
im[:, :, :] = 128
lim = np.zeros((width, height, 3), dtype='uint8')
if img.shape[0] >= img.shape[1]:
scale = img.shape[0] / height
new_width = int(img.shape[1] / scale)
diff = (width - new_width) // 2
img = cv2.resize(img, (new_width, height))
label_img = cv2.resize(label_img, (new_width, height))
im[:, diff:diff + new_width, :] = img
lim[:, diff:diff + new_width, :] = label_img
else:
scale = img.shape[1] / width
new_height = int(img.shape[0] / scale)
diff = (height - new_height) // 2
img = cv2.resize(img, (width, new_height))
label_img = cv2.resize(label_img, (width, new_height))
im[diff:diff + new_height, :, :] = img
lim[diff:diff + new_height, :, :] = label_img
lim = lim[:, :, 0]
seg_labels = np.zeros((height, width, nClasses))
for c in range(nClasses):
seg_labels[:, :, c] = (lim == c).astype(int)
im = np.float32(im) / 127.5 - 1
seg_labels = np.reshape(seg_labels, (width * height, nClasses))
x.append(im.transpose((2,0,1)))
y.append(seg_labels.transpose((1,0)))
return mx.nd.array(x), mx.nd.array(y)
class UnetDataIter(mx.io.DataIter):
def __init__(self, root_path, path_file, batch_size, n_classes, input_width, input_height, train=True):
f = open(path_file, 'r')
self.items = f.readlines()
f.close()
self._provide_data = [['data', (batch_size, 3, input_width, input_height)]]
self._provide_label = [['softmax_label', (batch_size, n_classes, input_width*input_height)]]
self.root_path = root_path
self.batch_size = batch_size
self.num_batches = len(self.items) // batch_size
self.n_classes = n_classes
self.input_height = input_height
self.input_width = input_width
self.train = train
self.reset()
def __iter__(self):
return self
def reset(self):
self.cur_batch = 0
self.shuffled_items = []
index = [n for n in range(len(self.items))]
if self.train:
random.shuffle(index)
for i in range(len(self.items)):
self.shuffled_items.append(self.items[index[i]])
def __next__(self):
return self.next()
@property
def provide_data(self):
return self._provide_data
@property
def provide_label(self):
return self._provide_label
def next(self):
if self.cur_batch == 0:
print("")
print("\r\033[k"+("Training " if self.train else "Validating ")+str(self.cur_batch)+"/"+str(self.num_batches), end=' ')
if self.cur_batch < self.num_batches:
data, label = get_batch(self.shuffled_items[self.cur_batch * self.batch_size:(self.cur_batch + 1) * self.batch_size], self.root_path, self.n_classes, self.input_height, self.input_width)
self.cur_batch += 1
return mx.io.DataBatch([data], [label])
else:
raise StopIteration
if __name__ =='__main__':
root_path = '/datasets/'
train_file = './data/seg_train.txt'
val_file = './data/seg_test.txt'
batch_size = 16
n_classes = 2
img_width = 256
img_height = 256
trainiter = UnetDataIter(root_path, train_file, batch_size, n_classes, img_width, img_height, True)
while True:
trainiter.next()
unet.py
import os
os.environ["MXNET_BACKWARD_DO_MIRROR"] = "1"
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
import mxnet as mx
from mxnet import ndarray as F
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from unetdataiter import UnetDataIter
import matplotlib.pyplot as plt
def dice_coef(y_true, y_pred):
intersection = mx.sym.sum(mx.sym.broadcast_mul(y_true, y_pred), axis=(1, 2, 3))
return mx.sym.broadcast_div((2. * intersection + 1.),(mx.sym.sum(y_true, axis=(1, 2, 3)) + mx.sym.sum(y_pred, axis=(1, 2, 3)) + 1.))
def dice_coef_loss(y_true, y_pred):
intersection = mx.sym.sum(mx.sym.broadcast_mul(y_true, y_pred), axis=1, )
return -mx.sym.broadcast_div((2. * intersection + 1.),(mx.sym.broadcast_add(mx.sym.sum(y_true, axis=1), mx.sym.sum(y_pred, axis=1)) + 1.))
def build_unet(batch_size, input_width, input_height, train=True):
data = mx.sym.Variable(name='data')
label = mx.sym.Variable(name='softmax_label')
# encode
# 256x256
conv1 = mx.sym.Convolution(data, num_filter=64, kernel=(3,3), pad=(1,1), name='conv1_1')
conv1 = mx.sym.BatchNorm(conv1, name='bn1_1')
conv1 = mx.sym.Activation(conv1, act_type='relu', name='relu1_1')
conv1 = mx.sym.Convolution(conv1, num_filter=64, kernel=(3,3), pad=(1,1), name='conv1_2')
conv1 = mx.sym.BatchNorm(conv1, name='bn1_2')
conv1 = mx.sym.Activation(conv1, act_type='relu', name='relu1_2')
pool1 = mx.sym.Pooling(conv1, kernel=(2,2), pool_type='max', name='pool1')
# 128x128
conv2 = mx.sym.Convolution(pool1, num_filter=128, kernel=(3,3), pad=(1,1), name='conv2_1')
conv2 = mx.sym.BatchNorm(conv2, name='bn2_1')
conv2 = mx.sym.Activation(conv2, act_type='relu', name='relu2_1')
conv2 = mx.sym.Convolution(conv2, num_filter=128, kernel=(3,3), pad=(1,1), name='conv2_2')
conv2 = mx.sym.BatchNorm(conv2, name='bn2_2')
conv2 = mx.sym.Activation(conv2, act_type='relu', name='relu2_2')
pool2 = mx.sym.Pooling(conv2, kernel=(2,2), pool_type='max', name='pool2')
# 64x64
conv3 = mx.sym.Convolution(pool2, num_filter=256, kernel=(3,3), pad=(1,1), name='conv3_1')
conv3 = mx.sym.BatchNorm(conv3, name='bn3_1')
conv3 = mx.sym.Activation(conv3, act_type='relu', name='relu3_1')
conv3 = mx.sym.Convolution(conv3, num_filter=256, kernel=(3,3), pad=(1,1), name='conv3_2')
conv3 = mx.sym.BatchNorm(conv3, name='bn3_2')
conv3 = mx.sym.Activation(conv3, act_type='relu', name='relu3_2')
pool3 = mx.sym.Pooling(conv3, kernel=(2,2), pool_type='max', name='pool3')
# 32x32
conv4 = mx.sym.Convolution(pool3, num_filter=256, kernel=(3,3), pad=(1,1), name='conv4_1')
conv4 = mx.sym.BatchNorm(conv4, name='bn4_1')
conv4 = mx.sym.Activation(conv4, act_type='relu', name='relu4_1')
conv4 = mx.sym.Convolution(conv4, num_filter=256, kernel=(3,3), pad=(1,1), name='conv4_2')
conv4 = mx.sym.BatchNorm(conv4, name='bn4_2')
conv4 = mx.sym.Activation(conv4, act_type='relu', name='relu4_2')
pool4 = mx.sym.Pooling(conv4, kernel=(2,2), pool_type='max', name='pool4')
# 16x16
conv5 = mx.sym.Convolution(pool4, num_filter=256, kernel=(3,3), pad=(1,1), name='conv5_1')
conv5 = mx.sym.BatchNorm(conv5, name='bn5_1')
conv5 = mx.sym.Activation(conv5, act_type='relu', name='relu5_1')
conv5 = mx.sym.Convolution(conv5, num_filter=256, kernel=(3,3), pad=(1,1), name='conv5_2')
conv5 = mx.sym.BatchNorm(conv5, name='bn5_2')
conv5 = mx.sym.Activation(conv5, act_type='relu', name='relu5_2')
pool5 = mx.sym.Pooling(conv5, kernel=(2,2), pool_type='max', name='pool5')
# 8x8
# decode
trans_conv6 = mx.sym.Deconvolution(pool5, num_filter=256, kernel=(2,2), stride=(1,1), no_bias=True, name='trans_conv6')
up6 = mx.sym.concat(*[trans_conv6, conv5], dim=1, name='concat6')
conv6 = mx.sym.Convolution(up6, num_filter=256, kernel=(3,3), pad=(1,1), name='conv6_1')
conv6 = mx.sym.BatchNorm(conv6, name='bn6_1')
conv6 = mx.sym.Activation(conv6, act_type='relu', name='relu6_1')
conv6 = mx.sym.Convolution(conv6, num_filter=256, kernel=(3,3), pad=(1,1), name='conv6_2')
conv6 = mx.sym.BatchNorm(conv6, name='bn6_2')
conv6 = mx.sym.Activation(conv6, act_type='relu', name='relu6_2')
trans_conv7 = mx.sym.Deconvolution(conv6, num_filter=256, kernel=(2,2), stride=(1,1), no_bias=True, name='trans_conv7')
up7 = mx.sym.concat(*[trans_conv7, conv4], dim=1, name='concat7')
conv7 = mx.sym.Convolution(up7, num_filter=256, kernel=(3,3), pad=(1,1), name='conv7_1')
conv7 = mx.sym.BatchNorm(conv7, name='bn7_1')
conv7 = mx.sym.Activation(conv7, act_type='relu', name='relu7_1')
conv7 = mx.sym.Convolution(conv7, num_filter=256, kernel=(3,3), pad=(1,1), name='conv7_2')
conv7 = mx.sym.BatchNorm(conv7, name='bn7_2')
conv7 = mx.sym.Activation(conv7, act_type='relu', name='relu7_2')
trans_conv8 = mx.sym.Deconvolution(conv7, num_filter=256, kernel=(2,2), stride=(1,1), no_bias=True, name='trans_conv8')
up8 = mx.sym.concat(*[trans_conv8, conv3], dim=1, name='concat8')
conv8 = mx.sym.Convolution(up8, num_filter=256, kernel=(3,3), pad=(1,1), name='conv8_1')
conv8 = mx.sym.BatchNorm(conv8, name='bn8_1')
conv8 = mx.sym.Activation(conv8, act_type='relu', name='relu8_1')
conv8 = mx.sym.Convolution(conv8, num_filter=256, kernel=(3,3), pad=(1,1), name='conv8_2')
conv8 = mx.sym.BatchNorm(conv8, name='bn8_2')
conv8 = mx.sym.Activation(conv8, act_type='relu', name='relu8_2')
trans_conv9 = mx.sym.Deconvolution(conv8, num_filter=128, kernel=(2,2), stride=(1,1), no_bias=True, name='trans_conv9')
up9 = mx.sym.concat(*[trans_conv9, conv2], dim=1, name='concat9')
conv9 = mx.sym.Convolution(up9, num_filter=128, kernel=(3,3), pad=(1,1), name='conv9_1')
conv9 = mx.sym.BatchNorm(conv9, name='bn9_1')
conv9 = mx.sym.Activation(conv9, act_type='relu', name='relu9_1')
conv9 = mx.sym.Convolution(conv9, num_filter=128, kernel=(3,3), pad=(1,1), name='conv9_2')
conv9 = mx.sym.BatchNorm(conv9, name='bn9_2')
conv9 = mx.sym.Activation(conv9, act_type='relu', name='relu9_2')
trans_conv10 = mx.sym.Deconvolution(conv9, num_filter=64, kernel=(2,2), stride=(1,1), no_bias=True, name='trans_conv10')
up10 = mx.sym.concat(*[trans_conv10, conv1], dim=1, name='concat10')
conv10 = mx.sym.Convolution(up10, num_filter=64, kernel=(3,3), pad=(1,1), name='conv10_1')
conv10 = mx.sym.BatchNorm(conv10, name='bn10_1')
conv10 = mx.sym.Activation(conv10, act_type='relu', name='relu10_1')
conv10 = mx.sym.Convolution(conv10, num_filter=64, kernel=(3,3), pad=(1,1), name='conv10_2')
conv10 = mx.sym.BatchNorm(conv10, name='bn10_2')
conv10 = mx.sym.Activation(conv10, act_type='relu', name='relu10_2')
###
conv11 = mx.sym.Convolution(conv10, num_filter=2, kernel=(1,1), name='conv11_1')
conv11 = mx.sym.sigmoid(conv11, name='softmax')
net = mx.sym.Reshape(conv11, (batch_size, 2, input_width*input_height))
if train:
loss = mx.sym.MakeLoss(dice_coef_loss(label, net), normalization='batch')
mask_output = mx.sym.BlockGrad(conv11, 'mask')
out = mx.sym.Group([loss, mask_output])
else:
# mask_output = mx.sym.BlockGrad(conv11, 'mask')
out = mx.sym.Group([conv11])
return outtrainunet.py
import os
os.environ["MXNET_BACKWARD_DO_MIRROR"] = "1"
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
import mxnet as mx
from mxnet import ndarray as F
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from unetdataiter import UnetDataIter
import matplotlib.pyplot as plt
from unet import build_unet
def main():
root_path = '../datasets/'
train_file = './data/seg_train.txt'
val_file = './data/seg_test.txt'
batch_size = 16
n_classes = 2
# img_width = 256
# img_height = 256
img_width = 96
img_height = 96
train_iter = UnetDataIter(root_path, train_file, batch_size, n_classes, img_width, img_height, True)
val_iter = UnetDataIter(root_path, val_file, batch_size, n_classes, img_width, img_height, False)
ctx = [mx.gpu(0)]
unet_sym = build_unet(batch_size, img_width, img_height)
unet = mx.mod.Module(unet_sym, context=ctx, data_names=('data',), label_names=('softmax_label',))
unet.bind(data_shapes=[['data', (batch_size, 3, img_width, img_height)]], label_shapes=[['softmax_label', (batch_size, n_classes, img_width*img_height)]])
unet.init_params(mx.initializer.Xavier(magnitude=6))
unet.init_optimizer(optimizer = 'adam',
optimizer_params=(
('learning_rate', 1E-4),
('beta1', 0.9),
('beta2', 0.99)
))
# unet.fit(train_iter, # train data
# eval_data=val_iter, # validation data
# #optimizer='sgd', # use SGD to train
# #optimizer_params={'learning_rate':0.1}, # use fixed learning rate
# eval_metric='acc', # report accuracy during training
# batch_end_callback = mx.callback.Speedometer(batch_size, 1), # output progress for each 100 data batches
# num_epoch=10) # train for at most 10 dataset passes
epochs = 20
smoothing_constant = .01
curr_losses = []
moving_losses = []
i = 0
best_val_loss = np.inf
for e in range(epochs):
while True:
try:
batch = next(train_iter)
except StopIteration:
train_iter.reset()
break
unet.forward_backward(batch)
loss = unet.get_outputs()[0]
unet.update()
curr_loss = F.mean(loss).asscalar()
curr_losses.append(curr_loss)
moving_loss = (curr_loss if ((i == 0) and (e == 0))
else (1 - smoothing_constant) * moving_loss + (smoothing_constant) * curr_loss)
moving_losses.append(moving_loss)
i += 1
val_losses = []
for batch in val_iter:
unet.forward(batch)
loss = unet.get_outputs()[0]
val_losses.append(F.mean(loss).asscalar())
val_iter.reset()
val_loss = np.mean(val_losses)
print("\nEpoch %i: Moving Training Loss %0.5f, Validation Loss %0.5f" % (e, moving_loss, val_loss))
unet.save_checkpoint('./unet_person_segmentation', e)
if __name__ =='__main__':
main()
以上是训练代码。
预测代码如下predict.py
import os
os.environ["MXNET_BACKWARD_DO_MIRROR"] = "1"
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
import sys
import cv2
import mxnet as mx
from mxnet import ndarray as F
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from unetdataiter import UnetDataIter
import matplotlib.pyplot as plt
from unet import build_unet
def post_process_mask(label, img_cols, img_rows, n_classes, p=0.5):
pr = label.reshape(n_classes, img_cols, img_rows).transpose([1,2,0]).argmax(axis=2)
return (pr*255).asnumpy()
def load_image(img, width, height):
im = np.zeros((height, width, 3), dtype='uint8')
im[:, :, :] = 128
if img.shape[0] >= img.shape[1]:
scale = img.shape[0] / height
new_width = int(img.shape[1] / scale)
diff = (width - new_width) // 2
img = cv2.resize(img, (new_width, height))
im[:, diff:diff + new_width, :] = img
else:
scale = img.shape[1] / width
new_height = int(img.shape[0] / scale)
diff = (height - new_height) // 2
img = cv2.resize(img, (width, new_height))
im[diff:diff + new_height, :, :] = img
im = np.float32(im) / 127.5 - 1
return [im.transpose((2,0,1))]
def main():
batch_size = 16
n_classes = 2
# img_width = 256
# img_height = 256
img_width = 96
img_height = 96
ctx = [mx.gpu(0)]
# sym, arg_params, aux_params = mx.model.load_checkpoint('unet_person_segmentation', 20)
# unet_sym = build_unet(batch_size, img_width, img_height, False)
# unet = mx.mod.Module(symbol=unet_sym, context=ctx, label_names=None)
sym, arg_params, aux_params = mx.model.load_checkpoint('unet_person_segmentation', 0)
unet = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
unet.bind(for_training=False, data_shapes=[['data', (batch_size, 3, img_width, img_height)]], label_shapes=unet._label_shapes)
unet.set_params(arg_params, aux_params, allow_missing=True)
testimg = cv2.imread(sys.argv[1], 1)
img = load_image(testimg, img_width, img_height)
unet.predict(mx.io.NDArrayIter(data=[img]))
outputs = unet.get_outputs()[0]
cv2.imshow('test', testimg)
cv2.imshow('mask', post_process_mask(outputs[0], img_width, img_height, n_classes))
cv2.waitKey()
if __name__ == '__main__':
if len(sys.argv) < 2:
print("illegal parameters")
sys.exit(0)
main()
剥离softmax保存参数用于ncnn模型转换，train2infer.py
import os
os.environ["MXNET_BACKWARD_DO_MIRROR"] = "1"
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
import sys
import cv2
import mxnet as mx
from mxnet import ndarray as F
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from unetdataiter import UnetDataIter
import matplotlib.pyplot as plt
from unet import build_unet
def main():
batch_size = 16
n_classes = 2
# img_width = 256
# img_height = 256
img_width = 96
img_height = 96
ctx = [mx.gpu(0)]
sym, arg_params, aux_params = mx.model.load_checkpoint(sys.argv[1], int(sys.argv[2]))
unet_sym = build_unet(batch_size, img_width, img_height, False)
unet = mx.mod.Module(symbol=unet_sym, context=ctx, label_names=None)
unet.bind(for_training=False, data_shapes=[['data', (batch_size, 3, img_width, img_height)]], label_shapes=unet._label_shapes)
unet.set_params(arg_params, aux_params, allow_missing=True)
unet.save_checkpoint('./unet_person_segmentation', 0)
if __name__ == '__main__':
if len(sys.argv) < 3:
print("illegal parameters")
sys.exit(0)
main()