起因
- 拿来小伙伴的代码,数据预处理阶段程序就跪了,找了找原因,内存用完了。要来他的top命令截图一看,呵呵,0.2t,小伙伴用的学校实验室的最好的机器,256GB的内存,玩起来当然没所谓,我用的公司的机器就略微寒酸了,内存32GB。我这个要玩的话,只能分批次读入图片,处理。
what i've done
- 认真找了找tensorflow的分批次读取数据的方案,有方法,找到了tf dataset的文档,功能强大,完全,可以一批一批高效的读进来数据给显卡时刻准备着。
- 后来还是放弃了tf的路子,现有的代码全是keras,把tf的代码嵌入进来,略微费劲儿了一点,况且我这次实验目标只是一个短期小实验,如果之后需要上线产品用再改用tf dataset或者TFRecord好了,在大量数据情况下配合集群hdfs会有更好的成效。
- 将来用的时候可以参考的这几个链接,描述tf读取数据的技巧:
- 找到了keras里有也有分批次读取数据的玩法,叫做fit_generator, 找到了官网文档,然后也找到了一位斯坦福同学的博客,我就按着这两个做下来,恩,完成了,效果不错。
"""stream data by batch"""
import numpy as np
import keras
from keras.utils import np_utils
from PIL import Image
class DataGenerator(keras.utils.Sequence):
def __init__(self, data, batch_size=128, dim=(32, 128, 3), n_classes=21000, shuffle=True):
"""
:param data: 包括了图片名称,路径,四个标签
:param batch_size: 一批训练128张图片
:param dim: 图片的维度
:param n_classes: 有多少种汉字
:param shuffle: 每个epoch末尾,是否打乱index顺序
"""
self.name_list, self.fullpath_list, self.y0, self.y1, self.y2, self.y3 = data
self.indexes = np.arange(len(self.name_list))
self.dim = dim
self.batch_size = batch_size
self.n_classes = n_classes
self.shuffle = shuffle
self.on_epoch_end()
def on_epoch_end(self):
"""
Updates indexes after each epoch
:return:
"""
if self.shuffle is True:
np.random.shuffle(self.indexes)
def __data_generation(self, indexes_temp):
"""
Generates data containing batch_size samples
:param indexes_temp:
:return:
"""
x = []
y = [[] for i in range(4)]
for index in indexes_temp:
img = Image.open(self.fullpath_list[index])
corp_img = img.crop((186, 0, 300, 30))
corp_img = corp_img.resize((128, 32))
image = np.array(corp_img)
image = image[..., :3]
x.append(image)
y[0].append(self.y0[index])
y[1].append(self.y1[index])
y[2].append(self.y2[index])
y[3].append(self.y3[index])
x = np.asarray(x, dtype=np.float32)
x = x / 255.0
y = [np.asarray(y[i], dtype=int) for i in range(4)]
y = [np_utils.to_categorical(y[i], self.n_classes) for i in range(4)]
return x, [y[0], y[1], y[2], y[3]]
def __len__(self):
"""
Denote the number of batches per epoch
:return:
"""
return int(np.floor(len(self.name_list) / self.batch_size))
def __getitem__(self, index_batch):
"""
Generate on batch of data
:param index_batch: the index of batch
:return:
"""
indexes_temp = self.indexes[index_batch * self.batch_size: (index_batch + 1) * self.batch_size]
x, [y0, y1, y2, y3] = self.__data_generation(indexes_temp)
return x, [y0, y1, y2, y3]
- ps小坑一个: 在validator用fit_generator的时候,tensorboard没法正常工作了,在github上看到https://github.com/keras-team/keras/issues/3358这还是一个open的issue,在博客下留言请教了,博客作者回复他也没有很好的解决方案。
what's more
- 数据可以分批次装下之后,可以正常运行了,不过又发现处理起来特别慢,显卡经常是一会儿100%然后就休息很长一段时间了,又仔细开始找原因,发现是IO效率低下了,说人话就是:大概是花了100秒钟读取图片到内存,然后CPU花了20秒钟裁切图片,然后送到GPU,GPU全力工作两秒钟就干完活了就在休息。
- improvement 1:
iostat -y 1
, 命令里看到io速度很慢,把数据全部挪到SSD硬盘上,tps
和KB_read/s
都是原来的10倍了你敢信,可能那台电脑上有不少人都在用机械硬盘工作导致那个盘速度太慢。 - improvement 2:读取图片和切图的python的库,从
cv2
换成了pillow-simd
, 从这里的benchmarks分析说后者在很多情况下比前者牛逼太多,换完后确实发现有提升。
result
看到 tps和KB_read/s的数据 很稳定了,持续的读进来图片, CPU使用率也很稳定,保持在17左右。然后间接带来的效果,GPU的利用率也就比原来高了不少,不会经常在休息了。
优化之后,在公司的电脑上训练内存消耗为之前的二十分之一,200GB->10GB,在我的显卡运算能力和显存都相对较小的情况下,我的模型训练时间没有增加,模型训练效果与之前相同。
future improvement
很明显可以多线程多个cpu同时开搞么,现阶段的改进已经足够这些天的实验了,就先不费功夫写多进程代码了。