記憶體溢位問題是參加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,train_batch_size,capacity)
X_valid_batch,y_valid_batch = inp.get_batch(X_valid,y_valid,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,batch_size=64,epochs=50,validation_data=(X_valid,y_valid),callbacks=[ckpt])
      
      del X_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),color_type)

獲取批次函式，其實就是一個generator

def get_train_batch(X_train,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],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,True),steps_per_epoch=1351,validation_data=get_train_batch(X_valid,False),validation_steps=52,callbacks=[ckpt,early_stop],max_queue_size=capacity,workers=1)

就是這麼簡單。但是當初從0到1的過程很難熬，每天都沒有進展，沒有頭緒，急躁佔據了思維的大部，熬過了這個階段，就會一切順利，不是運氣，而是踩過的從0到1的每個腳印累積的靈感的爆發，從0到1的腳印越多，後面的路越順利。

以上這篇完美解決TensorFlow和Keras大資料量記憶體溢位的問題就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。