模型加速--IO加速，tfRecord和keras Sequence

最新推荐文章于 2025-09-11 19:25:52 发布

原创

最新推荐文章于 2025-09-11 19:25:52 发布 · 925 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#tensorflow

本文介绍如何使用Keras的Sequence和TensorFlow的TFRecord优化大规模数据集的模型训练过程，减少数据读取时间，提高GPU利用率。

每次在训练模型时，尤其是训练数据较大时，都会大部分时间都会花在数据IO读写上，而不是真正的GPU计算，这也就意味着，GPU实际上很多时候是空闲等待状态！
在keras中可以通过sequence实现，在tensorflow中可以通过tfRecord实现。或者将图片以.npy的格式保存在本地，在训练的时候读取也会快很多。

如果将大规模数据一次性读进内存会很耗内存，可以使用tensorflow的queue和keras的sequence来存储数据。

1、keras的sequence
2、tensorflow的tfRecord

1、keras sequence

keras官方给出了参考实例，这里给了个常用的

import ast
import os
import numpy as np
import random
import math
from tensorflow.python.keras.preprocessing.image import img_to_array as img_to_array
from tensorflow.python.keras.preprocessing.image import load_img as load_img

def load_image(image_path, size):
    return img_to_array(load_img(image_path, target_size=(size, size))) / 255.

# shuffle好像有点问题,如有问题可参考https://www.kaggle.com/wrosinski/pretrained-cnn-albumentations
class KagglePlanetSequence(tf.keras.utils.Sequence):
    """
    在不把数据一次性读进内存的情况下，我们使用Sequence完成数据相对高效的IO
    """
    
    def __init__(self, df, data_path, im_size, batch_size, mode='train'):
        """
        df: pandas dataframe that contains columns with image names and labels
        data_path: path that contains the training images
        im_size: image size
        mode: when in training mode, data will be shuffled between epochs
        """
        self.df = df
        self.batch_size = batch_size
        self.im_size = im_size
        self.mode = mode
        
        # Take labels and a list of image locations in memory
        # ast.literal_eval类似eval，将字符转化为原有形式，更安全
        self.wlabels = self.df['weather_labels'].apply(lambda x: ast.literal_eval(x)).tolist()
        self.glabels = self.df['ground_labels'].apply(lambda x: ast.literal_eval(x)).tolist()
        self.image_list = self.df['image_name'].apply(lambda x: os.path.join(data_path, x + '.jpg')).tolist()

    def __len__(self):
        return int(math.ceil(len(self.df) / float(self.batch_size)))

    def on_epoch_end(self):
        # 每一轮之后对数据乱序
        self