完美解決TensorFlow和Keras大數據量內存溢出的問題

內存溢出問題是參加kaggle比賽或者做大數據量實驗的第一個攔路虎。

以前做的練手小項目導致新手產生一個慣性思維――讀取訓練集圖片的時候把所有圖讀到內存中，然後分批訓練。

其實這是有問題的，很容易導致OOM。現在內存一般16G，而訓練集圖片通常是上萬張，而且RGB圖，還很大，VGG16的圖片一般是224x224x3，上萬張圖片，16G內存根本不夠用。這時候又會想起――設置batch，但是那個batch的輸入參數卻又是圖片，它只是把傳進去的圖片分批送到顯卡，而我OOM的地方恰是那個“傳進去”的圖片，怎麼辦？

解決思路其實說來也簡單，打破思維定式就好了，不是把所有圖片讀到內存中，而是隻把所有圖片的路徑一次性讀到內存中。

大致的解決思路為：

將上萬張圖片的路徑一次性讀到內存中，自己實現一個分批讀取函數，在該函數中根據自己的內存情況設置讀取圖片，只把這一批圖片讀入內存中，然後交給模型，模型再對這一批圖片進行分批訓練，因為內存一般大於等於顯存，所以內存的批次大小和顯存的批次大小通常不相同。

下面代碼分別介紹Tensorflow和Keras分批將數據讀到內存中的關鍵函數。Tensorflow對初學者不太友好，所以我個人現階段更習慣用它的高層API Keras來做相關項目，下面的TF實現是之前不會用Keras分批讀時候參考的一些列資料，在模型訓練上仍使用Keras，只有分批讀取用了TF的API。

Tensorlow

在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的腳印越多，後面的路越順利。