【实验记录】U-Net训练自己的数据集（Keras）

论文地址：https://arxiv.org/abs/1505.04597
代码地址：https://github.com/zhixuhao/unet

环境配置

CUDA 8.0
cuDNN
Tensorflow 1.2.1
Keras 2.0.6
Python 3.5

# 创建虚拟环境
conda create -n u-net python=3.5
conda activate u-net
# 安装依赖
pip install tensorflow-gpu==1.2.1
pip install keras==2.0.6
pip install scikit-image
conda install numpy
conda install h5py

---

代码说明

可以运行一下代码中自带的数据集 membrane 看看：

python main.py

正常训练的话是这个样子：

main.py 是执行训练的主文件，其中：

• data_gen_args 定义了数据扩充操作
• trainGenerator 前四个参数分别为 1）训练的 batch_size，2）训练文件的主路径，3）训练图像的文件夹名称，4）训练图像对应标签的文件夹名称，接下来分别是 5）数据扩充操作，6、7）图像、标签的色彩模式，8、9）图像、标签的保存路径，10）是否多类别，11）类别个数（num_class 大于 2 的就属于多类别，flag 就应该设置为 True），12）是否保存，13）图片大小

---

训练 & 测试

1. 准备数据文件

在 ~/unet-master/data 下创建自己的数据集文件夹，比如 mydata。在 mydata 下创建 train 和 test 文件夹用于存放训练和测试数据。其中 train/image 和 train/label 中分别存放训练图像和对应的标签。test 下直接放所有测试图像即可。

2. 修改主文件 main.py

（1）修改训练文件路径及 batch size

myGene = trainGenerator(10,'data/mydata/train','image','label',data_gen_args,save_to_dir = None)

• 参数分别为 batch size（这里设为了 10），训练图像根目录，训练图像文件夹，训练图像对应标签文件夹

（2）修改 checkpoint 保存名称

model_checkpoint = ModelCheckpoint('unet_mydata.hdf5', monitor='loss',verbose=1, save_best_only=True)

（3）修改训练 epoch

model.fit_generator(myGene,steps_per_epoch=80,epochs=100,callbacks=[model_checkpoint],class_weight={1,245,245,245})

• steps_per_epoch 是每个 epoch 要迭代多少次，比如训练图像有 800 张，设训练 batch 为 10 的话，steps 就是 800 / 10 = 80；
• epochs 为训练的 epoch 数，这里设为了 100；

（4）修改测试路径

testGene = testGenerator("data/mydata/test")
# 测试图片有 400 张
results = model.predict_generator(testGene,400,verbose=1) 
saveResult("data/mydata/test",results)

3. 执行训练 & 测试

python main.py

另外，使用自己定义的 loss 函数，就在 model.py 修改使用的 loss 函数

model.compile(optimizer = Adam(lr = 1e-4), loss = 'ACLoss', metrics = ['accuracy'])

然后相应地，通过上面的 compile 找到了损失函数定义的 losses.py 文件，把自己的 loss 函数加进去就可以啦，输入参数为 (y_pred, y_true)

4. 多类别训练

源代码训练单类别，也就是只有前景（1）和背景（0）时没有问题，但是在进行多类别分割时报错：

ValueError: Error when checking target: expected conv2d_24 to have 4 dimensions, but got array with shape (5, 65536, 4)

原因：查了下原因是这个代码不适用于多类分割任务，参考了这篇文章，对代码做出修改。

解决：

（1）修改数据处理文件 data.py

def adjustData(img,mask,flag_multi_class,num_class):
    if(flag_multi_class):
        img = img / 255
        mask = mask[:,:,:,0] if(len(mask.shape) == 4) else mask[:,:,0]
        new_mask = np.zeros(mask.shape + (num_class,))
        for i in range(num_class):
            #for one pixel in the image, find the class in mask and convert it into one-hot vector
            #index = np.where(mask == i)
            #index_mask = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)
            #new_mask[index_mask] = 1
            new_mask[mask == i,i] = 1
        # 对这里进行了修改
        new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1],new_mask.shape[2],new_mask.shape[3])) if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))
        mask = new_mask
    elif(np.max(img) > 1):
        img = img / 255
        mask = mask /255
        mask[mask > 0.5] = 1
        mask[mask <= 0.5] = 0
    return (img,mask)

（2）修改模型文件 model.py 第大约第 53 行，因为报错中的 conv2d_24 实际就是 conv10，这里第一个参数 1 就是我们的类别数，将其修改为自己的类别数量即可。

# 将源代码：
conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)
# 修改为：
conv10 = Conv2D(4, 1, activation = 'sigmoid')(conv9)

（3）修改损失函数：源代码使用的是二元交叉熵，不适用于多类别分割问题，故这里改为 Dice loss。

# data.py line 39
new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1],new_mask.shape[2],new_mask.shape[3])) if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))

# model.py line 53
conv10 = Conv2D(4, 1, activation = 'sigmoid')(conv9)

---

报错记录与解决

【1】 ImportError: cannot import name 'tf_utils'

原因：keras 和 tensorflow 版本不兼容。
解决：tensorflow 1.2.1 和 keras 2.0.6 是 OK 的。

pip install tensorflow-gpu==1.2.1
pip install keras==2.0.6

=========================================================

【2】ImportError: `save_model` requires h5py

解决：安装 h5py：conda install h5py

=========================================================

【3】 tensorflow.python.framework.errors_impl.InternalError: Blas GEMM launch failed

原因：GPU 被占用
解决：确保一下 GPU 足够用呀~

=========================================================

【4】
Traceback (most recent call last):
  File "main.py", line 21, in 
    results = model.predict_generator(testGene,30,verbose=1)
  File "/data/zyy/usr/local/anaconda3/envs/u-net/lib/python3.5/site-packages/keras/legacy/interfaces.py", line 87, in wrapper
    return func(*args, **kwargs)
  File "/data/zyy/usr/local/anaconda3/envs/u-net/lib/python3.5/site-packages/keras/engine/training.py", line 2067, in predict_generator
    generator_output = next(output_generator)
StopIteration

Traceback (most recent call last):
  File "/data/zyy/usr/local/anaconda3/envs/u-net/lib/python3.5/threading.py", line 914, in _bootstrap_inner
    self.run()
  File "/data/zyy/usr/local/anaconda3/envs/u-net/lib/python3.5/threading.py", line 862, in run
    self._target(*self._args, **self._kwargs)
  File "/data/zyy/usr/local/anaconda3/envs/u-net/lib/python3.5/site-packages/keras/utils/data_utils.py", line 560, in data_generator_task
    generator_output = next(self._generator)
StopIteration

原因不知道，参考以下回答：
[1] https://github.com/zhixuhao/unet/issues/130
[2] https://stackoverflow.com/questions/46302911/what-raises-stopiteration-in-mine-keras-model-fit-generator

解决：

（1）根据 2067 行的报错，找到 training.py 代码中第 2003 行，将 max_queue_size 设置为 1

def predict_generator(self, generator, steps,
                          max_queue_size=1,    # modified
                          workers=1,
                          use_multiprocessing=False,
                          verbose=0):

=========================================================

【5】 Lossy conversion from float32 to uint8. Range [0, 1]. Convert image to uint8 prior to saving to suppress this warning.

原因：关于精度的警告，就是说从 float32 直接保存为 uint8 类型可能会损失精度。参考：https://www.jianshu.com/p/84b825b9e8a3
解决：将 image 转换为 uint8 类型。修改 data.py 最后的保存函数：

from skimage import img_as_ubyte

def saveResult(save_path,npyfile,flag_multi_class = False,num_class = 2):
    for i,item in enumerate(npyfile):
        img = labelVisualize(num_class,COLOR_DICT,item) if flag_multi_class else item[:,:,0]
        io.imsave(os.path.join(save_path,"%d_predict.png"%i),img_as_ubyte(img))  # modified

不过仍然会报一个警告，总比每张图片都有警告好

UserWarning: Possible precision loss when converting from float32 to uint8 .format(dtypeobj_in, dtypeobj_out))

=========================================================