tensorflowxun训练自己的数据集之从tfrecords读取数据

时间:2022-01-15 20:53:05

  当训练数据量较小时,采用直接读取文件的方式,当训练数据量非常大时,直接读取文件的方式太耗内存,这时应采用高效的读取方法,读取tfrecords文件,这其实是一种二进制文件。tensorflow为其内置了各种存储和读取的函数,方便调用。

  不知道为啥,从tfrecords中读取数据用于训练时,收敛得更快,更平稳。上面两个图是使用tfrecords的准确率和loss值变化,下面是直接读取文件的准确率和loss值变化。

tensorflowxun训练自己的数据集之从tfrecords读取数据tensorflowxun训练自己的数据集之从tfrecords读取数据

tensorflowxun训练自己的数据集之从tfrecords读取数据tensorflowxun训练自己的数据集之从tfrecords读取数据

1 生成记录样本的记录文件

 root_dir = os.getcwd()

 def getTrianList():
with open("train.txt","w") as f:
for file in os.listdir(root_dir+'\\dataSet'):
for picFile in os.listdir(root_dir+"\\dataSet\\"+file):
f.write("dataSet/"+file+"/"+picFile+" "+file+"\n")
print(picFile)
if __name__=="__main__":
getTrianList()

  将样本文件路径和标签统一记录到一个txt中,后面生成tfrecords文件就是通过读取这些信息。

  tensorflowxun训练自己的数据集之从tfrecords读取数据

  注意文件路径和标签之间采用空格,不要使用制表符。

2 读取txt存于数组中

 def load_file(example_list_file):
lines = np.genfromtxt(example_list_file,delimiter=" ",dtype=[('col1', 'S120'), ('col2', 'i8')])
examples = []
labels = []
for example,label in lines:
examples.append(example)
labels.append(label)
#convert to numpy array
return np.asarray(examples),np.asarray(labels),len(lines)

  这段代码主要用来读取第1步生成的txt,将文件路径和标签存于数组中

3 读取图片

 def extract_image(filename,height,width):
print(filename)
image = cv2.imread(filename)
image = cv2.resize(image,(height,width))
b,g,r = cv2.split(image)
rgb_image = cv2.merge([r,g,b])
return rgb_image

  使用cv2读取图片文件

4 转化为tfrecords文件

 def trans2tfRecord(trainFile,name,output_dir,height,width):
if not os.path.exists(output_dir) or os.path.isfile(output_dir):
os.makedirs(output_dir)
_examples,_labels,examples_num = load_file(train_file)
filename = name + '.tfrecords'
writer = tf.python_io.TFRecordWriter(filename)
for i,[example,label] in enumerate(zip(_examples,_labels)):
print("NO{}".format(i))
#need to convert the example(bytes) to utf-8
example = example.decode("UTF-8")
image = extract_image(example,height,width)
image_raw = image.tostring()
example = tf.train.Example(features=tf.train.Features(feature={
'image_raw':_bytes_feature(image_raw),
'height':_int64_feature(image.shape[0]),
'width': _int64_feature(32),
'depth': _int64_feature(32),
'label': _int64_feature(label)
}))
writer.write(example.SerializeToString())
writer.close()
 def _int64_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value])) def _bytes_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

5 从tfrecords中读取训练数据

 def read_tfRecord(file_tfRecord):
queue = tf.train.string_input_producer([file_tfRecord])
reader = tf.TFRecordReader()
_,serialized_example = reader.read(queue)
features = tf.parse_single_example(
serialized_example,
features={
'image_raw': tf.FixedLenFeature([], tf.string),
'height': tf.FixedLenFeature([], tf.int64),
'width':tf.FixedLenFeature([], tf.int64),
'depth': tf.FixedLenFeature([], tf.int64),
'label': tf.FixedLenFeature([], tf.int64)
}
)
image = tf.decode_raw(features['image_raw'],tf.uint8)
#height = tf.cast(features['height'], tf.int64)
#width = tf.cast(features['width'], tf.int64)
image = tf.reshape(image,[32,32,3])
image = tf.cast(image, tf.float32)
image = tf.image.per_image_standardization(image)
label = tf.cast(features['label'], tf.int64)
print(image,label)
return image,label

  从tfrecords文件中读取image和label,训练的时候,直接使用tf.train.batch函数生成用于训练的batch即可。

 image_batches,label_batches = tf.train.batch([image, label], batch_size=16, capacity=20)

  其余的部分跟之前的训练步骤一样。