问如何正确叠加capsnet (胶囊神经网络)？

EN

import os
from os.path import exists

import tensorflow as tf
import tensorflow_io as tfio

import matplotlib.pyplot as plt

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Variables
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
IMG_SHAPE = ( 32, 32, 1 ) 
tsk1_CLASSES_NUM = 2
initial_epochs = 5

PATH = os.path.join('F:\\datasets\\downloads\\Actors_jpg\\train\\Pikaploy', '*.jpg')
PATH_2 = os.path.join('F:\\datasets\\downloads\\Actors_jpg\\train\\Candidt Kibt', '*.jpg')
files = tf.data.Dataset.list_files(PATH)
files_2 = tf.data.Dataset.list_files(PATH_2)

list_file = []
list_file_actual = []
list_label = []
list_label_actual = [ 'Pikaploy', 'Pikaploy', 'Pikaploy', 'Pikaploy', 'Pikaploy', 'Candidt Kibt', 'Candidt Kibt', 'Candidt Kibt', 'Candidt Kibt', 'Candidt Kibt' ]
for file in files.take(5):
    image = tf.io.read_file( file )
    image =  tf.io.decode_jpeg( image, channels=0, ratio=1, fancy_upscaling=True, try_recover_truncated=False, acceptable_fraction=1, dct_method='INTEGER_FAST', name="decode_jpeg" )
    
    list_file_actual.append(image)
    image = tf.image.resize(image, [32,32], method='nearest')
    
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""
    : Image Masking
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""
    n_horizon = 64
    n_features = 16
    masking_layer = tf.keras.layers.Masking(mask_value=50, input_shape=(n_horizon, n_features))
    image = tf.constant( masking_layer(image)[:,:,0], shape=(32, 32, 1) )
    
    list_file.append(image)
    
    temp = tf.ones([ 20, 20 ]).numpy() * 0
    list_label.append( temp )
    
for file in files_2.take(5):
    image = tf.io.read_file( file )
    image =  tf.io.decode_jpeg( image, channels=0, ratio=1, fancy_upscaling=True, try_recover_truncated=False, acceptable_fraction=1, dct_method='INTEGER_FAST', name="decode_jpeg" )
    
    list_file_actual.append(image)
    image = tf.image.resize(image, [32,32], method='nearest')
    
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""
    : Image Masking
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""
    n_horizon = 64
    n_features = 16
    masking_layer = tf.keras.layers.Masking(mask_value=50, input_shape=(n_horizon, n_features))
    image = tf.constant( masking_layer(image)[:,:,0], shape=(32, 32, 1) )
    
    list_file.append(image)
    temp = tf.ones([ 20, 20 ]).numpy() * 9
    list_label.append( temp )
    
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)
    
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Definition / Class
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
def  build_model():

    branch_A_input = tf.keras.Input(shape=IMG_SHAPE)
    branch_A_rescale = tf.keras.layers.experimental.preprocessing.Rescaling(1./127.5, offset= -1)(branch_A_input)
    branch_A = tf.keras.layers.Dropout(0.3)(branch_A_rescale)
    branch_A = tf.keras.layers.Conv2D(filters = 128, kernel_size = 13, activation= 'swish', name = "base_conv_A")(branch_A)
    branch_A = tf.keras.layers.BatchNormalization(name = "base_batch_normalization_A")(branch_A)
    branch_A = tf.keras.Model(inputs=branch_A_input, outputs = branch_A)


    branch_B_input = tf.keras.Input(shape=IMG_SHAPE)
    branch_B_rescale = tf.keras.layers.experimental.preprocessing.Rescaling(1./127.5, offset= -1)(branch_B_input)
    branch_B = tf.keras.layers.Dropout(0.3)(branch_B_rescale)
    branch_B = tf.keras.layers.Conv2D(filters = 128, kernel_size = 13, activation= 'swish', name = "base_conv_B")(branch_B)
    branch_B = tf.keras.layers.BatchNormalization(name = "base_batch_normalization_B")(branch_B)
    branch_B = tf.keras.Model(inputs=branch_B_input, outputs = branch_B)

    merge = tf.keras.layers.Concatenate()([branch_A.output, branch_B.output])

    output_A = tf.keras.layers.Dense(tsk1_CLASSES_NUM, activation='softmax', name='4cls')(merge)
    output_B = tf.keras.layers.Dense(1, name='2cls')(merge)

    model = tf.keras.Model(inputs = [branch_A.input, branch_B.input] , outputs = [output_A, output_B], name="multi_task_model")

    optimizer = tf.keras.optimizers.get('adam')
    optimizer.learning_rate = 0.001
    losses = {'4cls': tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    '2cls': tf.keras.losses.BinaryCrossentropy(from_logits=True)}

    mtrcs = {
        "4cls": 'accuracy',
        "2cls": 'accuracy',
    }
    model.compile(optimizer=optimizer, loss= losses, metrics=mtrcs)

    return model
    
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: DataSet / Input
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
list_label = tf.constant( list_label, shape=( 10, 20, 20, 1 ) )
list_file = tf.constant( tf.cast( list_file, dtype=tf.int64), shape=( 10, 32, 32, 1 ) )

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Task
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
model = build_model()
history = model.fit( [list_file, list_file], list_label, epochs=initial_epochs)

input( '...' )

2022-12-02 09:39:30.019022: I tensorflow/stream_executor/cuda/cuda_dnn.cc:368] Loaded cuDNN version 8100
1/1 [==============================] - ETA: 0s - loss: nan - 4cls_loss: nan - 2cls_loss: 0.7376 - 4cls_accuracy: 0.2632 1/1 [==============================] - 2s 2s/step - loss: nan - 4cls_loss: nan - 2cls_loss: 0.7376 - 4cls_accuracy: 0.2632 - 2cls_accuracy: 0.3570
Epoch 2/5
1/1 [==============================] - ETA: 0s - loss: nan - 4cls_loss: nan - 2cls_loss: nan - 4cls_accuracy: 0.5000 - 21/1 [==============================] - 0s 11ms/step - loss: nan - 4cls_loss: nan - 2cls_loss: nan - 4cls_accuracy: 0.5000 - 2cls_accuracy: 0.5000