NVIDIA DLI 深度学习入门培训 | 特设三场！

4月28日/5月19日/5月26日一天密集式学习  快速带你入门阅读全文>

正文共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学院官网”：

www.leadai.org

请关注人工智能LeadAI公众号，查看更多专业文章

大家都在看

LSTM模型在问答系统中的应用

基于TensorFlow的神经网络解决用户流失概览问题

最全常见算法工程师面试题目整理（一）

最全常见算法工程师面试题目整理（二）

TensorFlow从1到2 | 第三章 深度学习革命的开端：卷积神经网络

装饰器 | Python高级编程

今天不如来复习下Python基础