【完全多线程的输入通道】MNIST图像分类

使用TensorFlow进行多线程MNIST数据输入与模型训练

最新推荐文章于 2024-04-29 11:07:13 发布

原创最新推荐文章于 2024-04-29 11:07:13 发布 · 185 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#tensorflow #队列 #python #深度学习 #图像识别

TensorFlow学习指南专栏收录该内容

8 篇文章

订阅专栏

该博客介绍了如何利用TensorFlow实现多线程的输入通道，结合MNIST数据集，从数据写入TFRecord文件，到数据加载、预处理，再到模型训练。通过创建QueueRunner和使用TFRecordReader，实现了数据的高效读取。最后，建立了一个简单的softmax分类模型并进行了训练。

部署运行你感兴趣的模型镜像

完全多线程的输入通道应用——MNIST手写数字图像分类

文章目录

完全多线程的输入通道应用——MNIST手写数字图像分类

简介

在本节我们将所有要介绍的内容放在一个 MNIST 图像的应用示例中，从数据写入到 TensorFlow 高效的文件格式，通过数据加载和预处理，到模型的训练。我们都用队列和多线程功能来实现这一点，并且依照这样的方式引入了一些更有用的组件来读取和处理 TensorFlow 中的数据。

项目设计

首先我们把 MNIST 数据写入TFRecord 中，使用之前的代码，此处不再过多赘述。详细内容可以参考这篇博客

	# Download data to save_dir
	data_sets = mnist.read_data_sets(save_dir, dtype=tf.uint8, reshape=False, validation_size=1000)
	
	data_splits = ["train", "test", "validation"]
	for d in range(len(data_splits)):
	    print("saving" + data_splits[d])
	    data_set = data_sets[d]
	
	    filename = os.path.join(save_dir, data_splits[d] + '.tfrecords')
	    writer = tf.python_io.TFRecordWriter(filename)
	
	    for index in range(data_set.images.shape[0]):
	        image = data_set.images[index].tostring()
	        example = tf.train.Example(features=tf.train.Features(feature={
	            'height' : tf.train.Feature(int64_list=tf.train.Int64List(value=[data_set.images.shape[1]])),
	            'width' : tf.train.Feature(int64_list=tf.train.Int64List(value=[data_set.images.shape[2]])),
	            'depth' : tf.train.Feature(int64_list=tf.train.Int64List(value=[data_set.images.shape[3]])),
	            'label' : tf.train.Feature(int64_list=tf.train.Int64List(value=[int(data_set.labels[index])])),
	            'image_raw' : tf.train.Feature(bytes_list=tf.train.BytesList(value=[image]))
	        }))
	
	        writer.write(example.SerializeToString())
	    writer.close()

使用 string_input_producer() 在幕后创建一个 QueueRunner , 将文件名子字符串输出到我们输入通道的队列中，这个文件名队列将被多个多线程共享

	filename = os.path.join(save_dir, "train.tfrecords")
	filename_queue = tf.train.string_input_producer([filename], num_epochs=10)

接着使用 TFReacordReader()从队列中读取文件，TFReacordReader()会以文件名队列为参数，按照文件名从 filename_queue 中出队
TFReacordReader()内部实现使用计算图的状态来跟踪将要被读取的 TFRecord 的位置，一块块地从磁盘的输入数据中加载

	reader = tf.TFRecordReader()
	_, serialized_example = reader.read(filename_queue)

返回一个 dict，解析数据并映射功能键到 Tensor 和 SparseTensor

	features = tf.parse_single_example(
	    serialized_example,
	    features={
	        'image_raw': tf.FixedLenFeature([], tf.string),
	        'label': tf.FixedLenFeature([], tf.int64)
	    })

解码原始字节字符串数据，进行基本的预处理，将像素值转换为浮点数，然后将图像实例进行重新排序，并通过 tf.train.shuffle_batch()
将其收集到 batch_size 批处理中，内部使用 RandomShuffleQueue 并累计示例，直到它包含 batch_size + min_sfter_dequeue个元素

	image = tf.decode_raw(features['image_raw'], tf.uint8)
	image.set_shape([784])
	image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
	label = tf.cast(features['label'], tf.int32)
	
	# 随机采集一批数据
	images_batch, labels_batch = tf.train.shuffle_batch(
	    [image, label], batch_size=16,
	    capacity=200, min_after_dequeue=100)

定义一个简单的 softmax 分类模型

	W = tf.get_variable("W", [28*28, 10])
	y_pred = tf.matmul(images_batch, W)
	loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y_pred, labels=labels_batch)
	
	loss_mean = tf.reduce_mean(loss)
	
	train_op = tf.train.AdamOptimizer().minimize(loss)
	
	config = tf.ConfigProto()
	config.gpu_options.allow_growth = True
	
	sess = tf.Session(config=config)
	init = tf.global_variables_initializer()
	sess.run(init)
	init = tf.local_variables_initializer()
	sess.run(init)

创建线程，进行训练

	coord = tf.train.Coordinator()
	threads = tf.train.start_queue_runners(sess=sess, coord=coord)
	print(threads)
	
	try:
	    step = 0
	    while not coord.should_stop() and step < 500:
	        step += 1
	        sess.run([train_op])
	        if step % 50 == 0:
	            loss_mean_val = sess.run(loss_mean)
	            print('step %d : %f' % (step, loss_mean_val))
	except tf.errors.OutOfRangeError:
	    print('Done training for %d epochs, %d steps.' % (NUM_EPOCHS, step))
	finally:
	    coord.request_stop()
	
	coord.join(threads)
	sess.close()

完整代码

	from __future__ import print_function
	import os
	import tensorflow as tf
	from tensorflow.contrib.learn.python.learn.datasets import mnist
	import numpy as np
	
	save_dir = 'c:/tmp/data'
	NUM_EPOCHS = 500
	
	# # Download data to save_dir
	# data_sets = mnist.read_data_sets(save_dir, dtype=tf.uint8, reshape=False, validation_size=1000)
	#
	# data_splits = ["train", "test", "validation"]
	# for d in range(len(data_splits)):
	#     print("saving" + data_splits[d])
	#     data_set = data_sets[d]
	#
	#     filename = os.path.join(save_dir, data_splits[d] + '.tfrecords')
	#     writer = tf.python_io.TFRecordWriter(filename)
	#
	#     for index in range(data_set.images.shape[0]):
	#         image = data_set.images[index].tostring()
	#         example = tf.train.Example(features=tf.train.Features(feature={
	#             'height' : tf.train.Feature(int64_list=tf.train.Int64List(value=[data_set.images.shape[1]])),
	#             'width' : tf.train.Feature(int64_list=tf.train.Int64List(value=[data_set.images.shape[2]])),
	#             'depth' : tf.train.Feature(int64_list=tf.train.Int64List(value=[data_set.images.shape[3]])),
	#             'label' : tf.train.Feature(int64_list=tf.train.Int64List(value=[int(data_set.labels[index])])),
	#             'image_raw' : tf.train.Feature(bytes_list=tf.train.BytesList(value=[image]))
	#         }))
	#
	#         writer.write(example.SerializeToString())
	#     writer.close()
	
	filename = os.path.join(save_dir, "train.tfrecords")
	# 在幕后创建一个QueueRunner, 将文件名子字符串输出到我们输入通道的队列中
	# 这个文件名队列将被多个多线程共享
	filename_queue = tf.train.string_input_producer([filename], num_epochs=10)
	
	# 接着使用 TFReacordReader()从队列中读取文件，TFReacordReader()会以文件名队列为参数，按照文件名从 filename_queue 中出队。
	# TFReacordReader()内部实现使用计算图的状态来跟踪将要被读取的 TFRecord 的位置，一块块地从磁盘的输入数据中加载
	reader = tf.TFRecordReader()
	_, serialized_example = reader.read(filename_queue)
	
	# 返回一个 dict，解析数据并映射功能键到 Tensor 和 SparseTensor
	features = tf.parse_single_example(
	    serialized_example,
	    features={
	        'image_raw': tf.FixedLenFeature([], tf.string),
	        'label': tf.FixedLenFeature([], tf.int64)
	    })
	
	# 解码原始字节字符串数据，进行基本的预处理，将像素值转换为浮点数，然后将图像实例进行重新排序，并通过 tf.train.shuffle_batch()
	# 将其收集到 batch_size 批处理中，内部使用 RandomShuffleQueue 并累计示例，直到它包含 batch_size + min_sfter_dequeue个元素
	image = tf.decode_raw(features['image_raw'], tf.uint8)
	image.set_shape([784])
	image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
	label = tf.cast(features['label'], tf.int32)
	
	# 随机采集一批数据
	images_batch, labels_batch = tf.train.shuffle_batch(
	    [image, label], batch_size=16,
	    capacity=200, min_after_dequeue=100)
	
	# 定义简单的 softmax 分类模型
	W = tf.get_variable("W", [28*28, 10])
	y_pred = tf.matmul(images_batch, W)
	loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y_pred, labels=labels_batch)
	
	loss_mean = tf.reduce_mean(loss)
	
	train_op = tf.train.AdamOptimizer().minimize(loss)
	
	config = tf.ConfigProto()
	config.gpu_options.allow_growth = True
	
	sess = tf.Session(config=config)
	init = tf.global_variables_initializer()
	sess.run(init)
	init = tf.local_variables_initializer()
	sess.run(init)
	
	coord = tf.train.Coordinator()
	threads = tf.train.start_queue_runners(sess=sess, coord=coord)
	print(threads)
	
	try:
	    step = 0
	    while not coord.should_stop() and step < 500:
	        step += 1
	        sess.run([train_op])
	        if step % 50 == 0:
	            loss_mean_val = sess.run(loss_mean)
	            print('step %d : %f' % (step, loss_mean_val))
	except tf.errors.OutOfRangeError:
	    print('Done training for %d epochs, %d steps.' % (NUM_EPOCHS, step))
	finally:
	    coord.request_stop()
	
	coord.join(threads)
	sess.close()