根据您的问题，确定网络的目的和输入数据响应。我创建了一个简单的自定义层，告诉您过程层只不过是简单的计算，每个层的数据过程输出来自卷积层和密集层。
示例：改进方法是通过输入/输出的比较，并试图扩大效果。
混淆矩阵，他试图看到整体模型训练的效果并用输入数据进行预测。

import os
from os.path import exists

import tensorflow as tf
import matplotlib.pyplot as plt

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
None
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
config = tf.config.experimental.set_memory_growth(physical_devices[0], True)
print(physical_devices)
print(config)

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Variables
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
BATCH_SIZE = 1
IMAGE_SIZE = ( 21, 16 ) 
objects_classes = [ 'plane', 'helicopter', 'truck' ]

checkpoint_path = "F:\\models\\checkpoint\\" + os.path.basename(__file__).split('.')[0] + "\\TF_DataSets_01.h5"
checkpoint_dir = os.path.dirname(checkpoint_path)

if not exists(checkpoint_dir) : 
    os.mkdir(checkpoint_dir)
    print("Create directory: " + checkpoint_dir)

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Class / Definition
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
class MyLSTMLayer( tf.keras.layers.LSTM ):
    def __init__(self, units, return_sequences, return_state):
        super(MyLSTMLayer, self).__init__( units, return_sequences=True, return_state=False )
        self.num_units = units

    def build(self, input_shape):
        self.kernel = self.add_weight("kernel",
        shape=[int(input_shape[-1]),
        self.num_units])

    def call(self, inputs):
        return tf.matmul(inputs, self.kernel)

def gen():
    train_generator = tf.keras.preprocessing.image.ImageDataGenerator(
        # rescale=1./255,
        # shear_range=0.2,
        # zoom_range=0.2,
        # horizontal_flip=True 
        )
    train_generator = train_generator.flow_from_directory(
        'F:\\temp\\image_catagorize',
        classes=[ 'plane', 'helicopter', 'truck' ],
        target_size=IMAGE_SIZE,
        batch_size=BATCH_SIZE,
        color_mode='grayscale',
        class_mode='sparse',    # None  # categorical   # binary    # sparse
        subset='training')

    return train_generator
    
train_generator = gen()
val_generator = train_generator

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Callback
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
class custom_callback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs={}):
        if( logs['accuracy'] >= 0.97 ):
            self.model.stop_training = True
    
custom_callback = custom_callback()

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Model Initialize
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
mycustomlayer = MyLSTMLayer( 64, True, False )

model = tf.keras.models.Sequential([
    tf.keras.layers.InputLayer(input_shape=( IMAGE_SIZE[0], IMAGE_SIZE[1], 1 ), name="Input_01Layer"),
    tf.keras.layers.UpSampling2D(size=(4, 4), name="UpSampling2DLayer_01"),
    tf.keras.layers.Normalization(mean=3., variance=2., name="NormalizationLayer_01"),
    tf.keras.layers.Normalization(mean=4., variance=6., name="NormalizationLayer_02"),
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', name="Conv2DLayer_01"),
    tf.keras.layers.MaxPooling2D((3, 3), name="MaxPooling2DLayer_01"),
    tf.keras.layers.Conv2D(32, (2, 2), activation='relu', name="Conv2DLayer_02"),
    tf.keras.layers.MaxPooling2D((2, 2), name="MaxPooling2DLayer_02"),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu', name="Conv2DLayer_03"),
    tf.keras.layers.Reshape(( 7 * 11, 64 )),
    ###
    mycustomlayer,
    ###
    tf.keras.layers.Dense(512, activation='relu', name="DenseLayer_01"),
    tf.keras.layers.Flatten(name="FlattenLayer_01"),
    tf.keras.layers.Dense(192, activation='relu', name="DenseLayer_02"),
    tf.keras.layers.Dense(3, name="DenseLayer_03"),
    
], name="MyModelClassification")

model.summary()

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Optimizer
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
optimizer = tf.keras.optimizers.SGD(
    learning_rate=0.000001,
    momentum=0.5,
    nesterov=True,
    name='SGD',
)

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Loss Fn
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""                               
lossfn = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True,
    reduction=tf.keras.losses.Reduction.AUTO,
    name='sparse_categorical_crossentropy'
)

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Model Summary
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
model.compile(optimizer=optimizer, loss=lossfn, metrics=['accuracy'])

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: FileWriter
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
if exists(checkpoint_path) :
    model.load_weights(checkpoint_path)
    print("model load: " + checkpoint_path)
    input("Press Any Key!")
    
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Training
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
history = model.fit(train_generator, validation_data=val_generator, batch_size=100, epochs=3, callbacks=[custom_callback] )
model.save_weights(checkpoint_path)

PATH = os.path.join('F:\\temp\\image_catagorize\\helicopter', '*.png')
files = tf.data.Dataset.list_files(PATH)
list_file = []
for file in files.take(20):
    image = tf.io.read_file( file )
    image = tf.io.decode_png( image, channels=1, dtype=tf.dtypes.uint8, name='decode_png' )
    image = tf.keras.preprocessing.image.img_to_array(image)
    image = tf.image.resize(image, IMAGE_SIZE, method='nearest')
    list_file.append(image)
    
PATH = os.path.join('F:\\temp\\image_catagorize\\plane', '*.png')
files = tf.data.Dataset.list_files(PATH)
for file in files.take(8):
    image = tf.io.read_file( file )
    image = tf.io.decode_png( image, channels=1, dtype=tf.dtypes.uint8, name='decode_png' )
    image = tf.keras.preprocessing.image.img_to_array(image)
    image = tf.image.resize(image, IMAGE_SIZE, method='nearest')
    list_file.append(image)
    
PATH = os.path.join('F:\\temp\\image_catagorize\\Truck', '*.png')
files = tf.data.Dataset.list_files(PATH)
for file in files.take(8):
    image = tf.io.read_file( file )
    image = tf.io.decode_png( image, channels=1, dtype=tf.dtypes.uint8, name='decode_png' )
    image = tf.keras.preprocessing.image.img_to_array(image)
    image = tf.image.resize(image, IMAGE_SIZE, method='nearest')
    list_file.append(image)
    
plt.figure(figsize=(6, 6))
plt.title("Actors recognitions")
for i in range(len(list_file)):
    img = tf.keras.preprocessing.image.array_to_img(
        list_file[i],
        data_format=None,
        scale=True
    )
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)
    predictions = model.predict(img_array)
    score = tf.nn.softmax(predictions[0])
    plt.subplot(6, 6, i + 1)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)
    plt.imshow(list_file[i])
    plt.xlabel(str(round(score[tf.math.argmax(score).numpy()].numpy(), 2)) + ":" +  str(objects_classes[tf.math.argmax(score)]))
    
plt.show()