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

 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

 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()

 # 读取图片函数 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)

 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)