tensorflow2装甲板id识别 2网络搭建与训练

tfrecord数据到可用数据集的转换

在上一篇文章中实现了tfrecord格式数据的读取
https://www.jianshu.com/p/88d09196bf07
但是读取的内容还不能直接被网络训练函数使用，因而需要对读到的数据进行简单处理

reader = tf.data.TFRecordDataset(record_path) # 打开一个TFrecord

读取数据后，图像数据进行解码和数据类型转化来适应网络计算的要求，最后将所得列表转化为tensor适配fit函数要求的格式，实现如下

def read_dataset(record_path):
    reader = tf.data.TFRecordDataset(record_path) # 打开一个TFrecord
    #reader = reader.shuffle (buffer_size = 1000) # 在缓冲区中随机打乱数据
    reader = reader.map (_parse_function) # 解析数据
    #for row in reader.take(1): #获取指定数量的数据集
    labels = []
    imgs = []
    for row in reader:                #遍历数据集
        label = tf.cast(row['label'],dtype=tf.float32)
        label = label - 1
        #此处应当注意tf.int8和tf.uint8的区别，使用错误将造成正常读入的图片解码结果与tfrecord解码结果不一致
        img = tf.io.decode_raw(row['img_raw'],out_type=tf.uint8)
        img = tf.cast(img,dtype=tf.float32)
        labels.append(label)
        imgs.append(img)

    np.random.seed(1024)
    np.random.shuffle(labels)
    np.random.seed(1024)
    np.random.shuffle(imgs)
    np.random.seed(1024)

    labels = tf.convert_to_tensor(labels)
    imgs =  tf.convert_to_tensor(imgs)
    return labels,imgs

应当注意tf.int8和tf.uint8的区别，使用错误将造成正常读入的图片解码结果与tfrecord解码结果不一致

训练过程与训练模型保存

该部分的内容与
https://www.jianshu.com/p/94cf2a32bbf0
中的差异并不大，这里直接贴出完整实现

import tensorflow as tf
import os
import numpy as np
#定义待解析数据集特征
feature_description = {
    'label': tf.io.FixedLenFeature([] , tf.int64, default_value=-1), # 默认值自己定义
    'img_raw' : tf.io.FixedLenFeature([], tf.string)
}
# 映射函数，用于解析一条example
def _parse_function (exam_proto): 
    return tf.io.parse_single_example (exam_proto, feature_description)
#读取返回数据集
def read_dataset(record_path):
    reader = tf.data.TFRecordDataset(record_path) # 打开一个TFrecord
    #reader = reader.shuffle (buffer_size = 1000) # 在缓冲区中随机打乱数据
    reader = reader.map (_parse_function) # 解析数据
    #for row in reader.take(1): #获取指定数量的数据集
    labels = []
    imgs = []
    for row in reader:                #遍历数据集
        label = tf.cast(row['label'],dtype=tf.float32)
        label = label - 1
        img = tf.io.decode_raw(row['img_raw'],out_type=tf.uint8)
        img = tf.cast(img,dtype=tf.float32)
        labels.append(label)
        imgs.append(img)

    np.random.seed(1024)
    np.random.shuffle(labels)
    np.random.seed(1024)
    np.random.shuffle(imgs)
    np.random.seed(1024)

    labels = tf.convert_to_tensor(labels)
    imgs =  tf.convert_to_tensor(imgs)
    return labels,imgs
      
if __name__ == '__main__':
    labels,imgs = read_dataset('./armor_train.tfrecords')
    #网络搭建
    model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(500,activation='relu',kernel_regularizer=tf.keras.regularizers.l2()),
        tf.keras.layers.Dense(128,activation='relu',kernel_regularizer=tf.keras.regularizers.l2()),
        tf.keras.layers.Dense(50,activation='relu',kernel_regularizer=tf.keras.regularizers.l2()),
        tf.keras.layers.Dense(8,activation='softmax',kernel_regularizer=tf.keras.regularizers.l2())
    ])
    #训练参数设置
    model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
        metrics=['sparse_categorical_accuracy']
    )
    #模型持久化设置
    ckpt_path = "./checkpoint/armor_id.ckpt"
    if(os.path.exists(ckpt_path + ".index")): #生成ckpt的同时会生成index文件,可通过该文件是否存在判断是否有预训练模型生成
        print("--load modle--")
        model.load_weights(ckpt_path)

    cp_callback = tf.keras.callbacks.ModelCheckpoint(
        filepath = ckpt_path,
        save_weights_only=True,                       #只保留模型参数
        save_best_only=True                                #只保留最优模型
    )
    #训练
    history = model.fit(imgs,labels,batch_size=32,epochs=50,validation_split=0.2,validation_freq=1,callbacks=[cp_callback])
    #网络结构和参数显示
    model.summary()

结果分析

训练100轮结果如下

curacy: 0.9874 - val_loss: 0.3584 - val_sparse_categorical_accuracy: 0.9764
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
flatten (Flatten)            multiple                  0         
_________________________________________________________________
dense (Dense)                multiple                  6144500   
_________________________________________________________________
dense_1 (Dense)              multiple                  64128     
_________________________________________________________________
dense_2 (Dense)              multiple                  6450      
_________________________________________________________________
dense_3 (Dense)              multiple                  408       
=================================================================
Total params: 6,215,486
Trainable params: 6,215,486
Non-trainable params: 0
_________________________________________________________________

后面还需要对网络的结构进行进一步的更改。不过可以确定的是训练流程已经跑通。