TensorFlow 数据读取模块调用过程（inception）

最新推荐文章于 2024-08-08 14:10:49 发布

__Sunny__

最新推荐文章于 2024-08-08 14:10:49 发布

阅读量1.7k

点赞数 1

分类专栏： TensorFlow 文章标签： TensorFlow

本文链接：https://blog.csdn.net/s_sunnyy/article/details/73529301

版权

本文介绍了TensorFlow中利用Inception模型处理多GPU数据的流程，包括数据分割、GPU Tower计算及参数同步。与cifar10不同，Inception将数据读取放在GPU Tower循环外部，仅创建单个数据输入管道，减少资源消耗。在imagenet_train.py中，数据读取、处理线程和样本队列的配置根据GPU数量调整，使用batch_join或shuffle_batch_join函数进行批处理。

摘要由CSDN通过智能技术生成

cifar10 的例子是针对单个 CPU/GPU 的，为考虑多个 GPU 的例子，选取了 TensorFlow 模型示例中 inception 的例子进行学习，个人感觉这个例子比 cifar10 多 GPU 的例子更好一点。

首先，inception 的处理思想：

1. 每一个 GPU 会用一批独立的数据计算梯度和估计值，将一大批数据分割成 split_batch_size 的数据分发到各个GPU上；
2. 在每个 GPU 上放置单独的模型副本，在单个模型拷贝中计算估计值和梯度的行为抽象到一个抽象对象为“tower”中；
3. 等所有 GPU 处理完一批数据后再同步更新模型的参数，在 CPU 上存储和更新所有模型的参数；

所以在 inception 这个例子中会有一个 batch_size 样本的分割的操作，并且与 cifar10 不同的是，它将读取数据的部分移到了 GPU Tower 循环的外部，这样就只创建了一个数据输入的管道，而不是每个GPU都创建一个数据输入管道。

（https://stackoverflow.com/questions/34273951/tensorflow-multi-gpu-single-input-queue）

下面具体看程序的部分：

models/inception/inception/imagenet_train.py

def main(_):
  dataset = ImagenetData(subset=FLAGS.subset)
  assert dataset.data_files()
  if tf.gfile.Exists(FLAGS.train_dir):
    tf.gfile.DeleteRecursively(FLAGS.train_dir)
  tf.gfile.MakeDirs(FLAGS.train_dir)
  inception_train.train(dataset) # ------------------->

models/inception/inception/inception_train.py

在 train 这个函数里中会根据 GPU 的个数重新设置数据处理线程的个数，并且在数据读取完成后对于多个 GPU 进行数据分割，然后将分割后的数据分发到各个 GPU Tower

def train(dataset):
  """Train on dataset for a number of steps."""
  with tf.Graph().as_default(), tf.device('/cpu:0'):
    ......
    # Get images and labels for ImageNet and split the batch across GPUs.
    assert FLAGS.batch_size % FLAGS.num_gpus == 0, (
        'Batch size must be divisible by number of GPUs')
    split_batch_size = int(FLAGS.batch_size / FLAGS.num_gpus)
	
    # 根据 GPU 的个数重新设置 num_preprocess_threads，使对应各个 GPU 都有一批预处理线程供给数据
    num_preprocess_threads = FLAGS.num_preprocess_threads * FLAGS.num_gpus
	
    # distorted_inputs() 读取并处理进行训练的数据
    images, labels = image_processing.distorted_inputs(dataset, num_preprocess_threads=num_preprocess_threads)  # -------------------->

    input_summaries = copy.copy(tf.get_collection(tf.GraphKeys.SUMMARIES))

    # Number of classes in the Dataset label set plus 1.
    # Label 0 is reserved for an (unused) background class.
    num_classes = dataset.num_classes() + 1
	
    # Split the batch of images and labels for towers.
    # 跨多个 GPUs 分割 batch 数据
    images_splits = tf.split(axis=0, num_or_size_splits=FLAGS.num_gpus, value=imag