如何用tf.data读取tfrecord数据4

使用tf.data数据来读取数据,代码如下

def read_and_decode(file_name,shuffle=True):
    def parser(record):
        features = tf.parse_single_example(record,
                                           features={
                                               'label': tf.FixedLenFeature([], tf.int64),
                                               'img_raw': tf.FixedLenFeature([], tf.string),
                                           })  # return image and label
        image = tf.decode_raw(features['img_raw'], tf.uint8)
        img = tf.reshape(image, [64, 64, 3])
        img = tf.cast(img, tf.float32)
        height = image_size
        width = image_size
        # Randomly crop a [height, width] section of the image.随机裁剪
        distorted_image = tf.random_crop(img, [height, width, 3])
        # Randomly flip the image horizontally.随机翻转
        distorted_image = tf.image.random_flip_left_right \
            (distorted_image)
        # 改变亮度
        distorted_image = tf.image.random_brightness(distorted_image,
                                                     max_delta=63)
        # 改变对比度
        distorted_image = tf.image.random_contrast(distorted_image,
                                                   lower=0.2, upper=1.8)
        img = tf.image.per_image_standardization(distorted_image)
        label = tf.cast(features['label'], tf.int32)
        return img,label

    dataset = tf.data.TFRecordDataset(file_name)
    if shuffle:
        dataset = dataset.map(parser).repeat().batch(batch_size).shuffle(buffer_size=1000)
    else:
        dataset = dataset.map(parser).repeat().batch(batch_size)

    iterator = dataset.make_one_shot_iterator()

    img_input, label = iterator.get_next()
    return img_input,label

其中parser函数是自定义的解码函数,之前的代码介绍怎么封装数据,现在需要parser来解码,map(parser)的作用是运行parser这个函数,tf.data.TFRecordDataset(file_name)中的file_name是这样的格式

file_name=["/home/python/dataset/train2.tfrecords","/home/python/dataset/train1.tfrecords"]

可以使用几个不同的tfrecords文件,调用的方式就是
images_train, labels_train = read_and_decode(file_name,shuffle=True)

,然后使用images_train, labels_train = sess.run(images_train,labels_train),然后就可以feed数据了,代码如下

_, loss_value = sess.run([train_op, total_loss],
                                             feed_dict={images_holder: images_batch,
                                                        labels_holder: labels_batch})
好了,这个就是完整的制作tfrecords文件,并且使用tf.data来调用的教程。

猜你喜欢

转载自blog.csdn.net/qwe2508/article/details/80680640