models.py

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
#from scipy.special import expit
from sklearn.metrics import silhouette_score
from tensorflow.keras.layers import *
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model
from tensorflow.keras.losses import mse
from tensorflow.keras.layers import MultiHeadAttention, LayerNormalization, Add
import os
import pandas as pd
#from glob import glob
from tensorflow.keras.regularizers import l2

def sampling(args):

    z_mean, z_log_var = args
    batch = K.shape(z_mean)[0]
    dim = K.int_shape(z_mean)[1]
    # by default, random_normal has mean=0 and std=1.0
    epsilon = K.random_normal(shape=(batch, dim))
    return z_mean + K.exp(0.5 * z_log_var) * epsilon


def get_MRI_VAE_3D(
    # input_shape=(119,143,103,1), 
    input_shape=(64,64,64,1), 
    latent_dim=2, 
    batch_size = 32, 
    disentangle=False,
    gamma=1,
    kernel_size = 3,
    filters = 32,
    intermediate_dim = 128,
    nlayers = 2,
    bias=True):

    image_size, _, _, channels = input_shape
    inputs = Input(shape=input_shape, name='encoder_input')
    x = inputs
    for i in range(nlayers):
        filters *= 2
        x = Conv3D(filters=filters,
                kernel_size=kernel_size,
                activation='relu',
                strides=2,
                padding='same',use_bias=bias)(x)

    # shape info needed to build decoder model
    shape = K.int_shape(--x)

    # generate latent vector Q(z|X)
    x = Flatten()(x)
    x = Dense(intermediate_dim, activation='relu',use_bias=bias)(x)
    z_mean = Dense(latent_dim, name='vae_mean',use_bias=bias)(x)
    z_log_var = Dense(latent_dim, name='vae_log_var',use_bias=bias)(x)

    # use reparameterization trick to push the sampling out as input
    # note that "output_shape" isn't necessary with the TensorFlow backend
    z = Lambda(sampling, output_shape=(latent_dim,), name='vae')([z_mean, z_log_var])

    # instantiate encoder model
    encoder = Model(inputs, [z_mean, z_log_var, z], name='encoder')

    # build decoder model
    latent_inputs = Input(shape=(latent_dim,), name='vae_sampling')
    x = Dense(intermediate_dim, activation='relu',use_bias=bias)(latent_inputs)
    x = Dense(shape[1] * shape[2] * shape[3] * shape[4], activation='relu',use_bias=bias)(x)
    x = Reshape((shape[1], shape[2], shape[3],shape[4]))(x)

    for i in range(nlayers):
        x = Conv3DTranspose(filters=filters,
                          kernel_size=kernel_size,
                          activation='relu',
                          strides=2,
                          padding='same',
                          use_bias=bias)(x)
        filters //= 2

    outputs = Conv3DTranspose(filters=1,
                            kernel_size=kernel_size,
                            activation='sigmoid',
                            padding='same',
                            name='decoder_output',
                            use_bias=bias)(x)

    # instantiate decoder model
    decoder = Model(latent_inputs, outputs, name='decoder')

    #     decoder.summary()

    # instantiate VAE model
    outputs = decoder(encoder(inputs)[2])
    vae = Model(inputs, outputs, name='vae')

    if disentangle:
        discriminator = Dense(1, activation='sigmoid')

        z1 = Lambda(lambda x: x[:int(batch_size/2),:int(latent_dim/2)])(z)
        z2 = Lambda(lambda x: x[int(batch_size/2):,:int(latent_dim/2)])(z)
        s1 = Lambda(lambda x: x[:int(batch_size/2),int(latent_dim/2):])(z)
        s2 = Lambda(lambda x: x[int(batch_size/2):,int(latent_dim/2):])(z)
        
        q_bar = tf.keras.layers.concatenate(
          [tf.keras.layers.concatenate([s1, z2], axis=1),
          tf.keras.layers.concatenate([s2, z1], axis=1)],
          axis=0)
        q = tf.keras.layers.concatenate(
          [tf.keras.layers.concatenate([s1, z1], axis=1),
          tf.keras.layers.concatenate([s2, z2], axis=1)],
          axis=0)

        q_bar_score = (discriminator(q_bar)+.1) *.85 # +.1 * .85 so that it's 0<x<1
        q_score = (discriminator(q)+.1) *.85 
        tc_loss = K.log(q_score / (1 - q_score)) 

        discriminator_loss = - K.log(q_score) - K.log(1 - q_bar_score)

    reconstruction_loss = mse(K.flatten(inputs), K.flatten(outputs))
    reconstruction_loss *= image_size * image_size


    kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
    kl_loss = K.sum(kl_loss, axis=-1)
    kl_loss *= -0.5
    if disentangle:
        vae_loss = K.mean(reconstruction_loss) + K.mean(kl_loss) + gamma * K.mean(tc_loss) + K.mean(discriminator_loss)
    else:
        vae_loss = K.mean(reconstruction_loss) + K.mean(kl_loss)

    vae.add_loss(vae_loss)
    opt = tf.keras.optimizers.Adam(learning_rate=0.001,beta_1=0.9,beta_2=0.999,epsilon=1e-07,amsgrad=False,name='Adam')
        
    vae.compile(optimizer=opt)
    

    if disentangle:
        vae.metrics_tensors = [reconstruction_loss, kl_loss, tc_loss, discriminator_loss]
    return encoder, decoder, vae



def DECODE_SZ(input_shape=(64,64,64,1), latent_dim=16, beta=1, disentangle=False, gamma=1, bias=True, batch_size = 64):


    image_size, _, _, channels = input_shape
    kernel_size = 3
    filters = 32
    intermediate_dim = 128
    nlayers = 2

    # build encoder model
    tg_inputs = Input(shape=input_shape, name='sz_inputs')
    bg_inputs = Input(shape=input_shape, name='cm_inputs')

    c_conv1 = Conv3D(filters=filters*2,
            kernel_size=kernel_size,
            activation='relu',
            strides=2,
            use_bias=bias,
            padding='same')

    c_conv2 = Conv3D(filters=filters*4,
            kernel_size=kernel_size,
            activation='relu',
            strides=2,
            use_bias=bias,
            padding='same')

    c_conv3 = Conv3D(filters=filters*8,
            kernel_size=kernel_size,
            activation='relu',
            strides=2,
            use_bias=bias,
            padding='same')
    
    c_bn = BatchNormalization()


    # generate latent vector Q(z|X)
    c_h_layer = Dense(intermediate_dim, activation='relu', use_bias=bias)
    c_mean_layer = Dense(latent_dim, name='c_mean', use_bias=bias)
    c_log_var_layer = Dense(latent_dim, name='c_log_var', use_bias=bias)
    c_layer = Lambda(sampling, output_shape=(latent_dim,), name='c')


    def c_encoder_func(inputs):
        c_h = inputs
        c_h = c_conv1(c_h)
        c_h = c_conv2(c_h)
        c_h = c_conv3(c_h)  
        c_h = c_bn(c_h)  
        shape = K.int_shape(c_h)
        c_h = Flatten()(c_h)
        c_h = c_h_layer(c_h)
        c_mean =  c_mean_layer(c_h)
        c_log_var =  c_log_var_layer(c_h)
        c = c_mean
        c = c_layer([c_mean, c_log_var])
        return c_mean, c_log_var, c, shape

    tg_c_mean, tg_c_log_var, tg_c, shape_c = c_encoder_func(tg_inputs)


########################################################################################################################

    s_conv1 = Conv3D(filters=filters*2,
            kernel_size=kernel_size,
            activation='relu',
            strides=2,
            use_bias=bias,
            padding='same')

    s_conv2 = Conv3D(filters=filters*4,
            kernel_size=kernel_size,
            activation='relu',
            strides=2,
            use_bias=bias,
            padding='same')
    s_conv3 = Conv3D(filters=filters*8,
            kernel_size=kernel_size,
            activation='relu',
            strides=2,
            use_bias=bias,
            padding='same')
    

    s_bn = BatchNormalization()

    
    # generate latent vector Q(z|X)
    s_h_layer = Dense(intermediate_dim, activation='relu', use_bias=bias)
    s_mean_layer = Dense(latent_dim, name='s_mean', use_bias=bias)
    s_log_var_layer = Dense(latent_dim, name='s_log_var', use_bias=bias)
    s_layer = Lambda(sampling, output_shape=(latent_dim,), name='s')

    def s_encoder_func(inputs):
        s_h = inputs
        s_h = s_conv1(s_h)
        s_h = s_conv2(s_h)

        s_h = s_conv3(s_h)
        s_h = s_bn(s_h)
        shape = K.int_shape(s_h)
        s_h = Flatten()(s_h)
        
        s_h = s_h_layer(s_h)
        s_mean =  s_mean_layer(s_h)
        s_log_var =  s_log_var_layer(s_h)
        s = s_mean
        s = s_layer([s_mean, s_log_var])
        return s_mean, s_log_var, s, shape

    tg_s_mean, tg_s_log_var, tg_s, shape_s = s_encoder_func(tg_inputs)
    bg_c_mean, bg_c_log_var, bg_c, _ = c_encoder_func(bg_inputs) 


      # instantiate encoder models
    c_encoder = tf.keras.models.Model(tg_inputs, [tg_c_mean, tg_c_log_var, tg_c], name='c_encoder')
    s_encoder = tf.keras.models.Model(tg_inputs, [tg_s_mean, tg_s_log_var, tg_s], name='s_encoder')

###############################################################################################################################
    
      # build decoder model
    latent_inputs = Input(shape=(2*latent_dim,), name='sc_sampling')

    x = Dense(intermediate_dim, activation='relu', use_bias=bias, kernel_regularizer=l2(0.01))(latent_inputs)
    x = Dense(shape_c[1] * shape_c[2] * shape_c[3] * shape_c[4], activation='relu', use_bias=bias)(x)
    x = Reshape((shape_c[1], shape_c[2], shape_c[3],shape_c[4]))(x)

    x = Conv3DTranspose(filters=filters*8,
                    kernel_size=kernel_size,
                    activation='relu',
                    strides=2,
                    use_bias=bias,
                    padding='same')(x)
    x = BatchNormalization()(x)  

    for i in range(nlayers):
        x = Conv3DTranspose(filters=filters,
                          kernel_size=kernel_size, 
                          activation='relu',
                          strides=2,
                          use_bias=bias,
                          padding='same')(x)
        x = BatchNormalization()(x)  
        filters //= 2

    outputs = Conv3DTranspose(filters=1,
                            kernel_size=kernel_size,
                            activation='sigmoid',
                            padding='same',
                            use_bias=bias,
                            name='decoder_output')(x)

    # instantiate decoder model
    decode_sz_decoder = Model(latent_inputs, outputs, name='decoder')

    def zeros_like(x):
        return tf.zeros_like(x)

    tg_outputs = decode_sz_decoder(tf.keras.layers.concatenate([tg_c, tg_s], -1))#患者数据时，同时走两个encoder
    zeros = tf.keras.layers.Lambda(zeros_like)(tg_c)

    bg_outputs = decode_sz_decoder(tf.keras.layers.concatenate([bg_c, zeros], -1))#正常人数据时，只走z_encoder，另外一部分用0向量代替

    decode_sz = tf.keras.models.Model(inputs=[tg_inputs, bg_inputs], 
                              outputs=[tg_outputs, bg_outputs], 
                              name='contrastive_vae')

    if disentangle:
        discriminator = Dense(1, activation='sigmoid')

        c1 = Lambda(lambda x: x[:int(batch_size/2),:])(tg_c)
        c2 = Lambda(lambda x: x[int(batch_size/2):,:])(tg_c)
        s1 = Lambda(lambda x: x[:int(batch_size/2),:])(tg_s)
        s2 = Lambda(lambda x: x[int(batch_size/2):,:])(tg_s)

        q_bar = tf.keras.layers.concatenate(
          [tf.keras.layers.concatenate([s1, c2], axis=1),
          tf.keras.layers.concatenate([s2, c1], axis=1)],
          axis=0)

        q = tf.keras.layers.concatenate(
          [tf.keras.layers.concatenate([s1, c1], axis=1),
          tf.keras.layers.concatenate([s2, c2], axis=1)],
          axis=0)

        q_bar_score = (discriminator(q_bar)+.1) *.85 
        q_score = (discriminator(q)+.1) *.85 
        tc_loss = K.log(q_score / (1 - q_score)) 
        discriminator_loss = - K.log(q_score) - K.log(1 - q_bar_score)
    else:
        tc_loss = 0
        discriminator_loss = 0


    reconstruction_loss = tf.keras.losses.mse(K.flatten(tg_inputs), K.flatten(tg_outputs)) 
    reconstruction_loss += tf.keras.losses.mse(K.flatten(bg_inputs), K.flatten(bg_outputs)) 
    reconstruction_loss *= input_shape[0] * input_shape[1] * input_shape[2] * input_shape[3]


    kl_loss = 1 + tg_c_log_var - tf.keras.backend.square(tg_c_mean) - tf.keras.backend.exp(tg_c_log_var)
    kl_loss += 1 + tg_s_log_var - tf.keras.backend.square(tg_s_mean) - tf.keras.backend.exp(tg_s_log_var)
    kl_loss += 1 + bg_c_log_var - tf.keras.backend.square(bg_c_mean) - tf.keras.backend.exp(bg_c_log_var)
    kl_loss = tf.keras.backend.sum(kl_loss, axis=-1)
    kl_loss *= -0.5
    
    decode_sz_loss = tf.keras.backend.mean(reconstruction_loss + beta*kl_loss + gamma*tc_loss + discriminator_loss)
    decode_sz.add_loss(decode_sz_loss)
    
    opt = tf.keras.optimizers.Adam(learning_rate=0.001,beta_1=0.9,beta_2=0.999,epsilon=1e-07,amsgrad=False,name='Adam')

    decode_sz.compile(optimizer=opt,run_eagerly=True)
    
    return decode_sz, c_encoder, s_encoder, decode_sz_decoder



# import tensorflow as tf
# from tensorflow.keras import Input, Model
# from tensorflow.keras.layers import (
#     Conv3D, Conv3DTranspose,
#     Dense, Flatten, Reshape, Lambda,
#     BatchNormalization, Cropping3D
# )
# from tensorflow.keras.regularizers import l2
# from tensorflow.keras import backend as K

# def sampling(args):
#     """Reparametrization trick"""
#     c_mean, c_log_var = args
#     batch = tf.shape(c_mean)[0]
#     dim = tf.shape(c_mean)[1]
#     epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
#     return c_mean + tf.exp(0.5 * c_log_var) * epsilon

# def DECODE_SZ_ORI(
#     # 修改 input_shape 为 (119, 143, 100, 1)
#     input_shape=(119, 143, 100, 1),
#     latent_dim=16,
#     beta=1,
#     disentangle=False,
#     gamma=1,
#     bias=True,
#     batch_size=64
# ):
#     depth, height, width, channels = input_shape
#     kernel_size = 3
#     filters = 32
#     intermediate_dim = 128
#     nlayers = 2

#     # -----------------------
#     # 构造 Encoder
#     # -----------------------
#     tg_inputs = Input(shape=input_shape, name='sz_inputs')
#     bg_inputs = Input(shape=input_shape, name='cm_inputs')

#     # 定义三层 Conv3D 和一个 BatchNormalization
#     c_conv1 = Conv3D(filters=filters * 2,
#                      kernel_size=kernel_size,
#                      activation='relu',
#                      strides=2,
#                      use_bias=bias,
#                      padding='same')
#     c_conv2 = Conv3D(filters=filters * 4,
#                      kernel_size=kernel_size,
#                      activation='relu',
#                      strides=2,
#                      use_bias=bias,
#                      padding='same')
#     c_conv3 = Conv3D(filters=filters * 8,
#                      kernel_size=kernel_size,
#                      activation='relu',
#                      strides=2,
#                      use_bias=bias,
#                      padding='same')
#     c_bn = BatchNormalization()

#     # 生成潜在变量所需的 Dense 层
#     c_h_layer = Dense(intermediate_dim, activation='relu', use_bias=bias)
#     c_mean_layer = Dense(latent_dim, name='c_mean', use_bias=bias)
#     c_log_var_layer = Dense(latent_dim, name='c_log_var', use_bias=bias)
#     c_layer = Lambda(sampling, output_shape=(latent_dim,), name='c')

#     def c_encoder_func(inputs):
#         c_h = inputs
#         c_h = c_conv1(c_h)   # (119 → 60, 143 → 72, 100 → 50) channels=64
#         c_h = c_conv2(c_h)   # (60 → 30, 72 → 36, 50 → 25)   channels=128
#         c_h = c_conv3(c_h)   # (30 → 15, 36 → 18, 25 → 13)   channels=256
#         c_h = c_bn(c_h)
#         shape = K.int_shape(c_h)  # (None, 15, 18, 13, 256)
#         c_h = Flatten()(c_h)
#         c_h = c_h_layer(c_h)
#         c_mean = c_mean_layer(c_h)
#         c_log_var = c_log_var_layer(c_h)
#         c = c_layer([c_mean, c_log_var])
#         return c_mean, c_log_var, c, shape

#     tg_c_mean, tg_c_log_var, tg_c, shape_c = c_encoder_func(tg_inputs)

#     # 分类器（“s” branch）的 Conv3D
#     s_conv1 = Conv3D(filters=filters * 2,
#                      kernel_size=kernel_size,
#                      activation='relu',
#                      strides=2,
#                      use_bias=bias,
#                      padding='same')
#     s_conv2 = Conv3D(filters=filters * 4,
#                      kernel_size=kernel_size,
#                      activation='relu',
#                      strides=2,
#                      use_bias=bias,
#                      padding='same')
#     s_conv3 = Conv3D(filters=filters * 8,
#                      kernel_size=kernel_size,
#                      activation='relu',
#                      strides=2,
#                      use_bias=bias,
#                      padding='same')
#     s_bn = BatchNormalization()

#     s_h_layer = Dense(intermediate_dim, activation='relu', use_bias=bias)
#     s_mean_layer = Dense(latent_dim, name='s_mean', use_bias=bias)
#     s_log_var_layer = Dense(latent_dim, name='s_log_var', use_bias=bias)
#     s_layer = Lambda(sampling, output_shape=(latent_dim,), name='s')

#     def s_encoder_func(inputs):
#         s_h = inputs
#         s_h = s_conv1(s_h)   # (119 → 60, 143 → 72, 100 → 50)
#         s_h = s_conv2(s_h)   # (60 → 30, 72 → 36, 50 → 25)
#         s_h = s_conv3(s_h)   # (30 → 15, 36 → 18, 25 → 13)
#         s_h = s_bn(s_h)
#         shape = K.int_shape(s_h)  # (None, 15, 18, 13, 256)
#         s_h = Flatten()(s_h)
#         s_h = s_h_layer(s_h)
#         s_mean = s_mean_layer(s_h)
#         s_log_var = s_log_var_layer(s_h)
#         s = s_layer([s_mean, s_log_var])
#         return s_mean, s_log_var, s, shape

#     tg_s_mean, tg_s_log_var, tg_s, shape_s = s_encoder_func(tg_inputs)
#     bg_c_mean, bg_c_log_var, bg_c, _ = c_encoder_func(bg_inputs)

#     # 最终创建 encoder 模型
#     c_encoder = Model(tg_inputs, [tg_c_mean, tg_c_log_var, tg_c], name='c_encoder')
#     s_encoder = Model(tg_inputs, [tg_s_mean, tg_s_log_var, tg_s], name='s_encoder')

#     # -----------------------
#     # 构造 Decoder
#     # -----------------------
#     latent_inputs = Input(shape=(2 * latent_dim,), name='sc_sampling')
#     x = Dense(intermediate_dim, activation='relu', use_bias=bias, kernel_regularizer=l2(0.01))(latent_inputs)

#     # 用 shape_c 构造一个全连接变形
#     d_c, h_c, w_c, f_c = shape_c[1], shape_c[2], shape_c[3], shape_c[4]
#     x = Dense(d_c * h_c * w_c * f_c, activation='relu', use_bias=bias)(x)
#     x = Reshape((d_c, h_c, w_c, f_c))(x)  # (15, 18, 13, 256)

#     # 与 encoder 中 Conv3D 层镜像，对应添加 Conv3DTranspose
#     x = Conv3DTranspose(filters=filters * 8,
#                         kernel_size=kernel_size,
#                         activation='relu',
#                         strides=2,
#                         use_bias=bias,
#                         padding='same')(x)
#     x = BatchNormalization()(x)
#     # 由 (15 → 30, 18 → 36, 13 → 26)

#     current_filters = filters  # 此时 filters=32
#     for i in range(nlayers):
#         x = Conv3DTranspose(filters=current_filters,
#                             kernel_size=kernel_size,
#                             activation='relu',
#                             strides=2,
#                             use_bias=bias,
#                             padding='same')(x)
#         x = BatchNormalization()(x)
#         current_filters //= 2
#     # 以上两层分别： (30 → 60, 36 → 72, 26 → 52)； (60 → 120, 72 → 144, 52 → 104)

#     outputs = Conv3DTranspose(filters=1,
#                               kernel_size=kernel_size,
#                               activation='sigmoid',
#                               padding='same',
#                               use_bias=bias,
#                               name='decoder_output')(x)
#     # 此时 shape = (120, 144, 104, 1)

#     # 用 Cropping3D 裁剪到 (119, 143, 100)
#     outputs = Cropping3D(cropping=((0, 1), (0, 1), (0, 4)), name='decoder_crop')(outputs)
#     # 最终 shape = (119, 143, 100, 1)

#     decode_sz_decoder = Model(latent_inputs, outputs, name='decoder')


#     tg_concat = tf.keras.layers.concatenate([tg_c, tg_s], axis=-1)
#     tg_outputs = decode_sz_decoder(tg_concat)


#     zeros = Lambda(lambda x: tf.zeros_like(x))(tg_c)
#     bg_concat = tf.keras.layers.concatenate([bg_c, zeros], axis=-1)
#     bg_outputs = decode_sz_decoder(bg_concat)

#     decode_sz = Model(inputs=[tg_inputs, bg_inputs],
#                  outputs=[tg_outputs, bg_outputs],
#                  name='contrastive_vae')

#     if disentangle:
#         discriminator = Dense(1, activation='sigmoid')
#         half_batch = batch_size // 2

#         c1 = Lambda(lambda x: x[:half_batch, :])(tg_c)
#         c2 = Lambda(lambda x: x[half_batch:, :])(tg_c)
#         s1 = Lambda(lambda x: x[:half_batch, :])(tg_s)
#         s2 = Lambda(lambda x: x[half_batch:, :])(tg_s)

#         q_bar = tf.keras.layers.concatenate(
#             [tf.keras.layers.concatenate([s1, c2], axis=1),
#              tf.keras.layers.concatenate([s2, c1], axis=1)],
#             axis=0
#         )
#         q = tf.keras.layers.concatenate(
#             [tf.keras.layers.concatenate([s1, c1], axis=1),
#              tf.keras.layers.concatenate([s2, c2], axis=1)],
#             axis=0
#         )

#         q_bar_score = (discriminator(q_bar) + 0.1) * 0.85
#         q_score = (discriminator(q) + 0.1) * 0.85
#         tc_loss = K.log(q_score / (1 - q_score))
#         discriminator_loss = -K.log(q_score) - K.log(1 - q_bar_score)
#     else:
#         tc_loss = 0
#         discriminator_loss = 0

#     # 重建损失：MSE(TG) + MSE(BG)，并乘以体素总数
#     reconstruction_loss = (
#         tf.keras.losses.mse(K.flatten(tg_inputs), K.flatten(tg_outputs)) +
#         tf.keras.losses.mse(K.flatten(bg_inputs), K.flatten(bg_outputs))
#     )
#     reconstruction_loss *= (depth * height * width * channels)

#     # KL 损失：tg_c, tg_s, bg_c 三个分支
#     kl_loss = (
#         1 + tg_c_log_var - K.square(tg_c_mean) - K.exp(tg_c_log_var) +
#         1 + tg_s_log_var - K.square(tg_s_mean) - K.exp(tg_s_log_var) +
#         1 + bg_c_log_var - K.square(bg_c_mean) - K.exp(bg_c_log_var)
#     )
#     kl_loss = K.sum(kl_loss, axis=-1) * -0.5

#     decode_sz_loss = K.mean(reconstruction_loss + beta * kl_loss + gamma * tc_loss + discriminator_loss)
#     decode_sz.add_loss(decode_sz_loss)

#     opt = tf.keras.optimizers.Adam(learning_rate=0.0005,
#                                    beta_1=0.9,
#                                    beta_2=0.999,
#                                    epsilon=1e-07,
#                                    amsgrad=False,
#                                    name='Adam')
#     decode_sz.compile(optimizer=opt, run_eagerly=True)

#     return decode_sz, c_encoder, s_encoder, decode_sz_decoder