完全多线程的输入通道应用——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()
本文示例参考《TensorFlow学习指南——深度学习系统构建详解》第八章第四节。
欢迎各位大佬交流讨论!