3D VAE的感知损失

如何解决3D VAE的感知损失

我正在为生物医学图像进行3D VAE实施。结果太模糊了，因此我正在寻求改善网络性能。许多人建议使用“感知损失”，但是我没有为此目的找到任何经过预先训练的3D-CNN。.我想知道是否存在其他方法来实现这种损失或改善性能的其他方法我的VAE网络。 我的代码：

class Sampling(tf.keras.layers.Layer):
    """Uses (z_mean,z_log_var) to sample z,the vector encoding a digit."""
    def call(self,inputs):
        z_mean,z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch,dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

def Encoder():
    inp = tf.keras.Input(shape=(32,256,1)) # prima era 64

    #enc = tf.keras.layers.Conv3D(8,(2,2,2),activation = 'relu',padding = 'same')(inp)
    #enc = tf.keras.layers.MaxPooling3D((2,padding = 'same')(enc)
 
    enc = tf.keras.layers.Conv3D(16,padding = 'same')(inp)
    enc = tf.keras.layers.MaxPooling3D((2,padding = 'same')(enc)
 
    enc = tf.keras.layers.Conv3D(32,padding = 'same')(enc)
    enc = tf.keras.layers.MaxPooling3D((2,padding = 'same')(enc)
 
    enc = tf.keras.layers.Conv3D(64,padding = 'same') (enc)

    enc = tf.keras.layers.Conv3D(32,padding = 'same') (enc)
    #enc = tf.keras.layers.Flatten()(enc)
    enc = tf.keras.layers.Conv3D(16,padding = 'same') (enc)
    '''
    # conv 2D 
    code = tf.keras.layers.Reshape((8,8,96)) (enc)
    code = tf.keras.layers.Conv2D(96,padding = 'same')(code)
    code = tf.keras.layers.MaxPooling2D((2,padding = 'same') (code)
    '''
    
    # latentent code vae
    latent_code = tf.keras.layers.Flatten()(enc)
    latent_code = tf.keras.layers.Dense(256,activation='relu')(latent_code)
    latent_mu = tf.keras.layers.Dense(32,activation='relu')(latent_code) # èprima era 10
    latent_sigma = tf.keras.layers.Dense(32,activation='relu')(latent_code) # prima era 10
    # Reparameterization trick
    #z = tf.keras.layers.Lambda(sample_z,output_shape=(128,),name='z')([latent_mu,latent_sigma])
    z = Sampling()([latent_mu,latent_sigma])
    encoder = tf.keras.Model(inp,[latent_mu,latent_sigma,z ],name = 'encoder')
    
    #encoder = tf.keras.Model(inp,enc)#[latent_mu,name = 'encoder')
    return encoder

def Decoder():
    z = tf.keras.Input(shape=(32,)) # prima era 10
    
    # start decoder
    rec = tf.keras.layers.Dense(1024,activation='relu')(z) # ripristino le dimensioni complete
    #rec = tf.keras.layers.BatchNormalization()(rec)
    rec = tf.keras.layers.Reshape((1,16))(rec) # riprestinate le dimensioni
    
    # traspose con2D
    code = tf.keras.layers.Conv3DTranspose(16,strides=(1,1,1),padding = 'same')(rec)
    #code = tf.keras.layers.BatchNormalization() (code)
    code = tf.keras.layers.UpSampling3D(size=(2,2))(code)
    # code = tf.keras.layers.Reshape((1,96)) (code)
    # end 2D

    
    #decoding
    dec = tf.keras.layers.Conv3DTranspose(32,activation='relu',padding='same')(code)
    #dec = tf.keras.layers.BatchNormalization()(dec)
    dec = tf.keras.layers.UpSampling3D(size=(2,2))(dec)
    dec = tf.keras.layers.Conv3DTranspose(64,padding='same')(dec)
    #dec = tf.keras.layers.BatchNormalization()(dec)
    dec = tf.keras.layers.UpSampling3D(size=(2,2))(dec)
    dec = tf.keras.layers.Conv3DTranspose(32,strides=1,2))(dec)
    dec = tf.keras.layers.Conv3DTranspose(16,2))(dec)
    #dec = tf.keras.layers.Conv3DTranspose(8,padding='same')(dec)
    #dec = tf.keras.layers.BatchNormalization()(dec)
    #dec = tf.keras.layers.UpSampling3D(size=(2,2))(dec)
    decoded = tf.keras.layers.Conv3D(1,(3,3,3),activation='sigmoid',padding='same')(dec)

    #model
    decoder = tf.keras.Model(inputs = z,outputs = decoded,name = 'decoder')
    return decoder


class ConvVAE3D(tf.keras.Model):
    def __init__(self,encoder,decoder,**kwargs):
        super(ConvVAE3D,self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
    
    def train_step(self,data):
        if isinstance(data,tuple):
            data = data[0]

        with tf.GradientTape() as tape:
            z_mean,z_log_var,z = self.encoder(data)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(data,reconstruction))#prima era binary crossentropy
            reconstruction_loss *= 256 * 256
            kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
            kl_loss = tf.reduce_mean(kl_loss)
            kl_loss *= -0.5
            total_loss = reconstruction_loss + kl_loss
            
        grads = tape.gradient(total_loss,self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads,self.trainable_weights))
        return {"loss": total_loss,"reconstruction_loss": reconstruction_loss,"kl_loss": kl_loss,}
    
    def test_step(self,tuple):
            data = data[0]

        z_mean,z = self.encoder(data)
        reconstruction = self.decoder(z)
        reconstruction_loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(data,reconstruction))
        reconstruction_loss *= 256 * 256
        kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
        kl_loss = tf.reduce_mean(kl_loss)
        kl_loss *= -0.5
        total_loss = reconstruction_loss + kl_loss
        return {
            "loss": total_loss,}
        
        
        
    def call(self,inputs): # implementa il forward pass
        z_mean,z = self.encoder(inputs)
        reconstruction = self.decoder(z)
        reconstruction_loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(inputs,reconstruction))
        reconstruction_loss *= 256 * 256
        kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
        kl_loss = tf.reduce_mean(kl_loss)
        kl_loss *= -0.5
        total_loss = reconstruction_loss + kl_loss
        self.add_metric(total_loss,name='loss',aggregation='mean')
        self.add_metric(reconstruction_loss,name='reconstruction_loss',aggregation='mean')
        self.add_metric(kl_loss,name='kl_loss',aggregation='mean')
        return reconstruction