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

日期: 2018-05
浏览: 212

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

  内存溢出问题是参加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的脚印越多,后面的路越顺利。


本文地址: http://www.61916.com/news_view_ 2841.html ,转载请注明出处。