TensorFlow和Keras解决大数据量内存溢出问题

正文共5771个字,1张图,预计阅读时间10分钟。

内存溢出问题是参加kaggle比赛或者做大数据量实验的第一个拦路虎。

以前做的练手小项目导致新手产生一个惯性思维——读取训练集图片的时候把所有图读到内存中,然后分批训练。

其实这是有问题的,很容易导致OOM。现在内存一般16G,而训练集图片通常是上万张,而且RGB图,还很大,VGG16的图片一般是224x224x3,上万张图片,16G内存根本不够用。这时候又会想起——设置batch,但是那个batch的输入参数却又是图片,它只是把传进去的图片分批送到显卡,而我OOM的地方恰是那个“传进去”的图片,怎么办?

解决思路其实说来也简单,打破思维定式就好了,不是把所有图片读到内存中,而是只把所有图片的路径一次性读到内存中。

大致的解决思路为:

将上万张图片的路径一次性读到内存中,自己实现一个分批读取函数,在该函数中根据自己的内存情况设置读取图片,只把这一批图片读入内存中,然后交给模型,模型再对这一批图片进行分批训练,因为内存一般大于等于显存,所以内存的批次大小和显存的批次大小通常不相同。

下面代码分别介绍Tensorflow和Keras分批将数据读到内存中的关键函数。Tensorflow对初学者不太友好,所以我个人现阶段更习惯用它的高层API Keras来做相关项目,下面的TF实现是之前不会用Keras分批读时候参考的一些列资料,在模型训练上仍使用Keras,只有分批读取用了TF的API。

TensorFlow

在input.py里写get_batch函数。

def get_batch(X_train, y_train, img_w, img_h, color_type, batch_size, capacity):
   '''
   Args:
       X_train: train img path list
       y_train: train labels list
       img_w: image width
       img_h: image height
       batch_size: batch size
       capacity: the maximum elements in queue
   Returns:
       X_train_batch: 4D tensor [batch_size, width, height, chanel],\
                       dtype=tf.float32
       y_train_batch: 1D tensor [batch_size], dtype=int32
   '''
   X_train = tf.cast(X_train, tf.string)

   y_train = tf.cast(y_train, tf.int32)    
   # make an input queue
   input_queue = tf.train.slice_input_producer([X_train, y_train])

   y_train = input_queue[1]
   X_train_contents = tf.read_file(input_queue[0])
   X_train = tf.image.decode_jpeg(X_train_contents, channels=color_type)

   X_train = tf.image.resize_images(X_train, [img_h, img_w], 
                                    tf.image.ResizeMethod.NEAREST_NEIGHBOR)

   X_train_batch, y_train_batch = tf.train.batch([X_train, y_train],
                                                 batch_size=batch_size,
                                                 num_threads=64,
                                                 capacity=capacity)

   y_train_batch = tf.one_hot(y_train_batch, 10)    return X_train_batch, y_train_batch
在train.py文件中训练(下面不是纯TF代码,model.fit是Keras的拟合,用纯TF的替换就好了)。
X_train_batch, y_train_batch = inp.get_batch(X_train, y_train, 
                                            img_w, img_h, color_type, 
                                            train_batch_size, capacity)
X_valid_batch, y_valid_batch = inp.get_batch(X_valid, y_valid, 
                                            img_w, img_h, color_type, 
                                            valid_batch_size, capacity)with tf.Session() as sess:

   coord = tf.train.Coordinator()
   threads = tf.train.start_queue_runners(coord=coord)   
 try:       
 for step in np.arange(max_step):            
if coord.should_stop() :                
break
           X_train, y_train = sess.run([X_train_batch, 
                                            y_train_batch])
           X_valid, y_valid = sess.run([X_valid_batch,
                                            y_valid_batch])
 
           ckpt_path = 'log/weights-{val_loss:.4f}.hdf5'
           ckpt = tf.keras.callbacks.ModelCheckpoint(ckpt_path, 
                                                     monitor='val_loss', 
                                                     verbose=1, 
                                                     save_best_only=True, 
                                                     mode='min')
           model.fit(X_train, y_train, batch_size=64, 
                         epochs=50, verbose=1,
                         validation_data=(X_valid, y_valid),
                         callbacks=[ckpt])            
           del X_train, y_train, X_valid, y_valid    
except tf.errors.OutOfRangeError:
       print('done!')    finally:
       coord.request_stop()
   coord.join(threads)

   sess.close()

Keras

keras文档中对fit、predict、evaluate这些函数都有一个generator,这个generator就是解决分批问题的。

关键函数:fit_generator

# 读取图片函数

def get_im_cv2(paths, img_rows, img_cols, color_type=1, normalize=True):
   '''
   参数:
       paths:要读取的图片路径列表
       img_rows:图片行
       img_cols:图片列
       color_type:图片颜色通道
   返回: 
       imgs: 图片数组
   '''
   # Load as grayscale
   imgs = []    for path in paths:        
if color_type == 1:
           img = cv2.imread(path, 0)        
elif color_type == 3:
           img = cv2.imread(path)        
# Reduce size
       resized = cv2.resize(img, (img_cols, img_rows))       
 if normalize:
           resized = resized.astype('float32')
           resized /= 127.5
           resized -= 1. 
 
       imgs.append(resized)        

   return np.array(imgs).reshape(len(paths), img_rows, img_cols, color_type)

获取批次函数,其实就是一个generator

def get_train_batch(X_train, y_train, batch_size, img_w, img_h, color_type, is_argumentation):
   '''
   参数:
       X_train:所有图片路径列表
       y_train: 所有图片对应的标签列表
       batch_size:批次
       img_w:图片宽
       img_h:图片高
       color_type:图片类型
       is_argumentation:是否需要数据增强
   返回: 
       一个generator,
x: 获取的批次图片 
y: 获取的图片对应的标签
   '''
   while 1:        
for i in range(0, len(X_train), batch_size):
           x = get_im_cv2(X_train[i:i+batch_size], img_w, img_h, color_type)
           y = y_train[i:i+batch_size]            
if is_argumentation:                
# 数据增强
               x, y = img_augmentation(x, y)            

# 最重要的就是这个yield,它代表返回,返回以后循环还是会继续,然后再返回。就比如有一个机器一直在作累加运算,但是会把每次累加中间结果告诉你一样,直到把所有数加完

yield({'input': x}, {'output': y})

训练函数

result = model.fit_generator(generator=get_train_batch(X_train, y_train, train_batch_size, img_w, img_h, color_type, True), 
         steps_per_epoch=1351, 
         epochs=50, verbose=1,
         validation_data=get_train_batch(X_valid, y_valid, valid_batch_size,img_w, img_h, color_type, False),
         validation_steps=52,
         callbacks=[ckpt, early_stop],
         max_queue_size=capacity,

         workers=1)

就是这么简单。但是当初从0到1的过程很难熬,每天都没有进展,没有头绪,急躁占据了思维的大部,熬过了这个阶段,就会一切顺利,不是运气,而是踩过的从0到1的每个脚印累积的灵感的爆发,从0到1的脚印越多,后面的路越顺利。

原文链接:https://www.jianshu.com/p/5bdae9dcfc9c

原文发布于微信公众号 - 人工智能LeadAI(atleadai)

原文发表时间:2018-04-18

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏TensorFlow从0到N

TensorFlow从0到1 - 12 - TensorFlow构建3层NN玩转MNIST

上一篇 11 74行Python实现手写体数字识别展示了74行Python代码完成MNIST手写体数字识别,识别率轻松达到95%。这算不上一个好成绩,不过我并...

5175
来自专栏京东技术

【解题报告】看雪·京东2018CTF—京东AI CTF大挑战特别题

近日,京东安全联合安全界的黄埔军校看雪论坛举办了一次线上CTF大赛,近3w人参赛。参与解答“京东AI CTF大挑战特别题”的同学有1435人,最终解出题目的只有...

1632
来自专栏一心无二用,本人只专注于基础图像算法的实现与优化。

肤色检测算法 - 基于二次多项式混合模型的肤色检测。

   由于能力有限,算法层面的东西自己去创新的很少,很多都是从现有的论文中学习,然后实践的。       本文涉及的很多算法,在网络上也有不少同类型的文章,但是...

31611
来自专栏大学生计算机视觉学习DeepLearning

手指静脉细化算法过程原理解析 以及python实现细化算法

1874
来自专栏云计算教程系列

如何使用Scikit-learn在Python中构建机器学习分类器

机器学习是计算机科学、人工智能和统计学的研究领域。机器学习的重点是训练算法以学习模式并根据数据进行预测。机器学习特别有价值,因为它让我们可以使用计算机来自动化决...

2265
来自专栏CDA数据分析师

教你一招 | Python实现无向图最短路径

一心想学习算法,很少去真正静下心来去研究,前几天趁着周末去了解了最短路径的资料,用python写了一个最短路径算法。算法是基于带权无向图去寻找两个点之间的最短路...

8175
来自专栏大学生计算机视觉学习DeepLearning

手指静脉细化算法过程原理解析 以及python实现细化算法

4865
来自专栏深度学习入门与实践

【深度学习系列】PaddlePaddle之数据预处理

  上篇文章讲了卷积神经网络的基本知识,本来这篇文章准备继续深入讲CNN的相关知识和手写CNN,但是有很多同学跟我发邮件或私信问我关于PaddlePaddle如...

2668
来自专栏技术墨客

TensorFlow入门 原

本文将初步向码农和程序媛们介绍如何使用TensorFlow进行编程。在阅读之前请先 安装TensorFlow,此外为了能够更好的理解本文的内容,阅读之前需要了解...

1152
来自专栏人工智能LeadAI

使用TensorFlow实现手写识别(Softmax)

准备工作 由于将TensorFlow安装到了Conda的tensorflow环境,虽然可以用Jupyter notebook打开,但是没有提示,写代码不方便,所...

4035

扫码关注云+社区

领取腾讯云代金券