深度学习12-TFRecord详解

1.TFRecord简介

1）TFRecord是什么

TFRecord --> Example --> feature --> key-value键值对，并且value的取值有三种

2）为什么用TFRecord

为什么要用TFRecord？

3）TFRecord格式

2.写入TFRecord

3.读取TFRecord文件

4.案例实战-猫狗图片分类

import tensorflow as tf
import os

# 处理文件路径
data_dir = 'datasets'

train_cats_dir = data_dir + '/train/cats/'
train_dogs_dir = data_dir + '/train/dogs/'

train_tfrecord_file = data_dir + '/train/train.tfrecords'

test_cats_dir = data_dir + '/valid/cats/'
test_dogs_dir = data_dir + '/valid/dogs/'

test_tfrecord_file = data_dir + '/valid/test.tfrecords'

# 将数据存储为TFRecord文件

# 文件名字变成list
train_cat_filenames = [train_cats_dir + filename for filename in os.listdir(train_cats_dir)]
train_dog_filenames = [train_dogs_dir + filename for filename in os.listdir(train_dogs_dir)]
train_filenames = train_cat_filenames + train_dog_filenames

# 将猫类的标签设为0，dog类的标签设为1
train_labels = [0] * len(train_cat_filenames) + [1] * len(train_dog_filenames)

# 写入TFRcord文件
with tf.io.TFRecordWriter(train_tfrecord_file) as writer:
    for filename, label in zip(train_filenames,train_labels):
        # 读取数据集图片到内存，image为一个byte类型的字符串
        image = open(filename,'rb').read()
        
        # 建立tf.train.Feature字典
        feature = {
            'image' : tf.train.Feature(bytes_list = tf.train.BytesList(value=[image])),    # 图片是一个Bytes对象
            'label' : tf.train.Feature(int64_list = tf.train.Int64List(value=[label]))     # 标签是一个Int对象
        }
        
        # 通过字典建立Example
        example = tf.train.Example(features=tf.train.Features(feature=feature))
        # 将Example序列化
        serialized = example.SerializeToString()
        # 写入TFRecord文件
        writer.write(serialized)
        
test_cat_filenames = [test_cats_dir + filename for filename in os.listdir(test_cats_dir)]
test_dog_filenames = [test_dogs_dir + filename for filename in os.listdir(test_dogs_dir)]
test_filenames = test_cat_filenames + test_dog_filenames

# 将猫类的标签设为0，dog类的标签设为1
test_labels = [0] * len(test_cat_filenames) + [1] * len(test_dog_filenames)  
with tf.io.TFRecordWriter(test_tfrecord_file) as writer:
    for filename, label in zip(test_filenames, test_labels):
        image = open(filename, 'rb').read()     # 读取数据集图片到内存，image 为一个 Byte 类型的字符串
        feature = {                             # 建立 tf.train.Feature 字典
            'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),  # 图片是一个 Bytes 对象
            'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))   # 标签是一个 Int 对象
        }
        example = tf.train.Example(features=tf.train.Features(feature=feature)) # 通过字典建立 Example
        serialized = example.SerializeToString() #将Example序列化
        writer.write(serialized)   # 写入 TFRecord 文件

# 读取TFRecord文件

# 定义Feature结构，告诉解码器每个Feature的类型是什么
feature_description = {
    'image': tf.io.FixedLenFeature([],tf.string),
    'label' : tf.io.FixedLenFeature([],tf.int64),
}

# 定义解码函数
def _parse_example(example_string):
    # 将TFRecord文件中的每一个序列化的tf.train.Example解码
    feature_dict = tf.io.parse_single_example(example_string, feature_description)
    # 解码JPEG图片
    feature_dict['image'] = tf.io.decode_jpeg(feature_dict['image'])   
    # 处理大小与像素
    feature_dict['image'] = tf.image.resize(feature_dict['image'],[256, 256]) / 255.0
    return feature_dict['image'], feature_dict['label']


batch_size = 32

# 读取TFRecord文件
train_dataset = tf.data.TFRecordDataset(train_tfrecord_file)

# 解码
train_dataset = train_dataset.map(_parse_example)
for image,label in train_dataset.take(1):
    print(image.shape,label)
# (256, 256, 3) tf.Tensor(0, shape=(), dtype=int64)

# 模型批量读取
train_dataset = train_dataset.shuffle(buffer_size = 23000)
train_dataset = train_dataset.batch(batch_size)
# (32,256,256,3)
# 优化
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)

# 测试集的读取
test_dataset = tf.data.TFRecordDataset(test_tfrecord_file)
test_dataset = test_dataset.map(_parse_example)
test_dataset = test_dataset.batch(batch_size)


# 定义CNN模型
class CNNModel(tf.keras.models.Model):
    def __init__(self):
        super(CNNModel,self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(32,3,activation='relu')
        self.maxpool1 = tf.keras.layers.MaxPooling2D()
        self.conv2 = tf.keras.layers.Conv2D(32,5,activation='relu')
        self.maxpool2 = tf.keras.layers.MaxPooling2D()
        self.flatten = tf.keras.layers.Flatten()
        self.d1 = tf.keras.layers.Dense(64,activation='relu')
        self.d2 = tf.keras.layers.Dense(2,activation='softmax')
        
    def call(self,x):
        # 定义前向传播
        x = self.conv1(x)
        x = self.maxpool1(x)
        x = self.conv2(x)
        x = self.maxpool2(x)
        x = self.flatten(x)
        x = self.d1(x)
        x = self.d2(x)
        
        return x

learning_rate = 0.001
model = CNNModel()

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam(learning_rate)

# 损失与评估
train_loss = tf.keras.metrics.Mean(name = 'train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name = 'test_accuracy')

# batch
# 将动态图转换为静态图，静态图执行效率高
@tf.function
def train_step(images,labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_object(labels,predictions)
    # 计算梯度
    gradients = tape.gradient(loss,model.trainable_variables)
    optimizer.apply_gradients(zip(gradients,model.trainable_variables))
    
    train_loss(loss)
    train_accuracy(labels,predictions)  # update
    
@tf.function
def test_step(images,labels):
    predictions = model(images)
    t_loss = loss_object(labels,predictions)
    
    test_loss(t_loss)
    test_accuracy(labels,predictions)

# 模型训练
EPOCHS = 10

for epoch in range(EPOCHS):
    # 重置评估指标
    train_loss.reset_states()
    train_accuracy.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()
    
    for images,labels in train_dataset:
        train_step(images,labels)
        
    for test_images,test_labels in test_dataset:
        test_step(images,labels)
        
    template = 'Epoch {}, Loss : {}, Accuracy : {},Test Loss : {},Test Accuracy : {}'
    
    # 打印
    print(template.format(epoch + 1,
                         train_loss.result(),
                         train_accuracy.result() * 100,
                         test_loss.result(),
                         test_accuracy.result() * 100
                         ))