我正在为生物医学图像进行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,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,2,2), padding = 'same')(enc)
enc = tf.keras.layers.Conv3D(16, (2,2,2), activation = 'relu', padding = 'same')(inp)
enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same')(enc)
enc = tf.keras.layers.Conv3D(32, (2,2,2), activation = 'relu', padding = 'same')(enc)
enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same')(enc)
enc = tf.keras.layers.Conv3D(64, (2,2,2), activation = 'relu', padding = 'same')(enc)
enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same') (enc)
enc = tf.keras.layers.Conv3D(32, (2,2,2), activation = 'relu', padding = 'same')(enc)
enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same') (enc)
#enc = tf.keras.layers.Flatten()(enc)
enc = tf.keras.layers.Conv3D(16, (2,2,2), activation = 'relu', padding = 'same')(enc)
enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same') (enc)
'''
# conv 2D
code = tf.keras.layers.Reshape((8,8,96)) (enc)
code = tf.keras.layers.Conv2D(96,(2,2), activation = 'relu', padding = 'same')(code)
code = tf.keras.layers.MaxPooling2D((2,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, latent_sigma, z ], 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, 8, 8, 16))(rec) # riprestinate le dimensioni
# traspose con2D
code = tf.keras.layers.Conv3DTranspose(16,(2,2,2), strides=(1,1,1),activation = 'relu', padding = 'same')(rec)
#code = tf.keras.layers.BatchNormalization() (code)
code = tf.keras.layers.UpSampling3D(size=(2, 2, 2))(code)
# code = tf.keras.layers.Reshape((1,8,8,96)) (code)
# end 2D
#decoding
dec = tf.keras.layers.Conv3DTranspose(32, (2,2,2), strides=(1,1,1) , activation='relu', padding='same')(code)
#dec = tf.keras.layers.BatchNormalization()(dec)
dec = tf.keras.layers.UpSampling3D(size=(2, 2, 2))(dec)
dec = tf.keras.layers.Conv3DTranspose(64, (2,2,2), strides=(1,1,1) , activation='relu', padding='same')(dec)
#dec = tf.keras.layers.BatchNormalization()(dec)
dec = tf.keras.layers.UpSampling3D(size=(2, 2, 2))(dec)
dec = tf.keras.layers.Conv3DTranspose(32, (2,2,2), strides=1, activation='relu', padding='same')(dec)
#dec = tf.keras.layers.BatchNormalization()(dec)
dec = tf.keras.layers.UpSampling3D(size=(2, 2, 2))(dec)
dec = tf.keras.layers.Conv3DTranspose(16, (2,2,2), strides=1, activation='relu', padding='same')(dec)
#dec = tf.keras.layers.BatchNormalization()(dec)
dec = tf.keras.layers.UpSampling3D(size=(2, 2, 2))(dec)
#dec = tf.keras.layers.Conv3DTranspose(8, (2,2,2), strides=1, activation='relu', padding='same')(dec)
#dec = tf.keras.layers.BatchNormalization()(dec)
#dec = tf.keras.layers.UpSampling3D(size=(2, 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, data):
if isinstance(data, tuple):
data = data[0]
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))
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,
"reconstruction_loss": reconstruction_loss,
"kl_loss": kl_loss,}
def call(self, inputs): # implementa il forward pass
z_mean, z_log_var, 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