tensorflow2.x一个完整的分类任务训练流程

import os
import tensorflow as tf
import glob
from model import AlexNet_v1

def make_dataset(image_dir, batch_size):
    class_dic = dict(zip(os.listdir(image_dir+'/train/'), range(len(os.listdir(image_dir+'/train/')))))
    iniv_class_dic = dict(zip(range(len(os.listdir(image_dir+'/train/'))), os.listdir(image_dir+'/train/')))

    train_image_ls = glob.glob(image_dir+'/train/*/*.jpg')
    train_label_ls = [class_dic[file.split('/')[-2]] for file in train_image_ls]

    val_image_ls = glob.glob(image_dir+'/val/*/*.jpg')
    val_label_ls = [class_dic[file.split('/')[-2]] for file in val_image_ls]

    AUTOTUNE = tf.data.experimental.AUTOTUNE
    def fn (imagefile, label):
        image = tf.io.read_file(imagefile)
        image = tf.image.decode_jpeg(image)
        image <