USER
import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Reshape, Conv2DTranspose, Conv2D, Flatten
from tensorflow.keras.models import Model
from tensorflow.keras.datasets import fashion_mnist
import numpy as np
from tensorflow.keras.layers import Layer
from tensorflow.keras.datasets import mnist
# Załadowanie danych MNIST
(x_train_mnist, _), (x_test_mnist, _) = mnist.load_data()
x_train_mnist = x_train_mnist.astype('float32') / 255.
x_test_mnist = x_test_mnist.astype('float32') / 255.
x_train_mnist = np.expand_dims(x_train_mnist, -1) # Dodanie kanału
x_test_mnist = np.expand_dims(x_test_mnist, -1)
from tensorflow.keras.datasets import fashion_mnist
# Załadowanie danych Fashion MNIST
(x_train_fashion, _), (x_test_fashion, _) = fashion_mnist.load_data()
x_train_fashion = x_train_fashion.astype('float32') / 255.
x_test_fashion = x_test_fashion.astype('float32') / 255.
x_train_fashion = np.expand_dims(x_train_fashion, -1) # Dodanie kanału: (60000, 28, 28, 1)
x_test_fashion = np.expand_dims(x_test_fashion, -1)
# Sprawdzenie kształtu
print(f"x_train_fashion shape: {x_train_fashion.shape}") # Powinno być (60000, 28, 28, 1)
print(f"x_test_fashion shape: {x_test_fashion.shape}") # Powinno być (10000, 28, 28, 1)
x_train_fashion shape: (60000, 28, 28, 1)
x_test_fashion shape: (10000, 28, 28, 1)
x_train_cifar shape: (50000, 32, 32, 3)
x_test_cifar shape: (10000, 32, 32, 3)
import os
import zipfile
import gdown # Biblioteka do pobierania plików z Google Drive
# Bezpośredni link do pliku ZIP
file_id = '1A2dNWabg6_um-V3lhw1tyead5hCpjaW8'
zip_url = f'https://drive.google.com/uc?id={file_id}'
# Ścieżka do pobranego pliku ZIP i lokalizacja rozpakowania
zip_path = '/content/image.zip'
local_extract_path = '/content/images'
# Pobieranie pliku ZIP
if not os.path.exists(zip_path): # Sprawdzanie, czy plik ZIP już istnieje
print(f"Pobieranie pliku ZIP z Google Drive: {zip_url}...")
gdown.download(zip_url, zip_path, quiet=False)
else:
print(f"Plik ZIP już istnieje: {zip_path}")
# Rozpakowanie archiwum
if not os.path.exists(local_extract_path): # Sprawdzanie, czy katalog już istnieje
os.makedirs(local_extract_path) # Tworzenie katalogu, jeśli nie istnieje
print(f"Rozpakowywanie pliku ZIP: {zip_path} do {local_extract_path}...")
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(local_extract_path)
print("Rozpakowywanie zakończone!")
else:
print(f"Katalog docelowy już istnieje: {local_extract_path}")
import glob
import os
import zipfile # Dodano moduł zipfile
from sklearn.model_selection import train_test_split
import numpy as np
from tensorflow.keras.preprocessing.image import load_img, img_to_array
# Path to CelebA images
celeba_path = '/content/images/image' # Update this path as needed
def load_celeba_images(path, img_size=(64, 64)):
image_paths = glob.glob(os.path.join(path, '*.jpg'))
images = []
for img_path in image_paths:
img = load_img(img_path, target_size=img_size)
img = img_to_array(img) / 255.
images.append(img)
return np.array(images)
# Load CelebA data
x_celeba = load_celeba_images(celeba_path)
print(f"x_celeba shape: {x_celeba.shape}") # e.g., (30000, 64, 64, 3)
# Split into training and validation sets
x_train_celeba, x_val_celeba = train_test_split(x_celeba, test_size=0.3, random_state=42)
x_celeba shape: (30000, 64, 64, 3)
class Sampling(Layer):
def call(self, inputs):
z_mean, z_log_var = inputs
epsilon = tf.keras.backend.random_normal(shape=tf.shape(z_mean))
return z_mean + tf.exp(0.5 * z_log_var) * epsilon
def build_conv_encoder(input_shape, latent_dim):
input_img = Input(shape=input_shape)
x = Conv2D(32, kernel_size=3, strides=2, activation='relu', padding='same')(input_img)
x = Conv2D(64, kernel_size=3, strides=2, activation='relu', padding='same')(x)
x_shape = tf.keras.backend.int_shape(x)[1:] # Zachowanie kształtu przed spłaszczeniem
x = Flatten()(x)
x = Dense(128, activation='relu')(x)
z_mean = Dense(latent_dim, name='z_mean')(x)
z_log_var = Dense(latent_dim, name='z_log_var')(x)
z = Sampling()([z_mean, z_log_var])
encoder = Model(input_img, [z_mean, z_log_var, z], name='encoder')
return encoder, x_shape
def build_conv_decoder(x_shape, latent_dim, output_shape):
latent_inputs = Input(shape=(latent_dim,))
x = Dense(np.prod(x_shape), activation='relu')(latent_inputs)
x = Reshape(x_shape)(x)
x = Conv2DTranspose(64, kernel_size=3, strides=2, activation='relu', padding='same')(x)
x = Conv2DTranspose(32, kernel_size=3, strides=2, activation='relu', padding='same')(x)
# Calculate current size after upsampling
current_size = x_shape[0] * 2**2 # Adjust based on strides
target_size = output_shape[0]
while current_size < target_size:
x = Conv2DTranspose(16, kernel_size=3, strides=1, activation='relu', padding='same')(x)
current_size += 1 # Adjust as needed
# Remove `break` to allow multiple layers if needed
outputs = Conv2DTranspose(
output_shape[2],
kernel_size=3,
activation='sigmoid',
padding='same'
)(x)
decoder = Model(latent_inputs, outputs, name='decoder')
return decoder
class VAE(Model):
def __init__(self, encoder, decoder, kl_weight=1.0, **kwargs):
super(VAE, self).__init__(**kwargs)
self.encoder = encoder
self.decoder = decoder
self.kl_weight = kl_weight
self.reconstruction_loss_tracker = tf.keras.metrics.Mean(name="reconstruction_loss")
self.kl_loss_tracker = tf.keras.metrics.Mean(name="kl_loss")
def call(self, inputs):
z_mean, z_log_var, z = self.encoder(inputs)
reconstruction = self.decoder(z)
return reconstruction
def train_step(self, data):
with tf.GradientTape() as tape:
z_mean, z_log_var, z = self.encoder(data)
reconstruction = self.decoder(z)
# Flatten the images
data_flat = tf.reshape(data, [tf.shape(data)[0], -1])
reconstruction_flat = tf.reshape(reconstruction, [tf.shape(reconstruction)[0], -1])
# Compute per-sample reconstruction loss
loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=False, reduction=tf.keras.losses.Reduction.NONE)
per_sample_loss = loss_fn(data_flat, reconstruction_flat)
reconstruction_loss = tf.reduce_mean(per_sample_loss)
# KL divergence loss
kl_loss = -0.5 * tf.reduce_mean(
tf.reduce_sum(
1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var),
axis=1
)
)
total_loss = reconstruction_loss + self.kl_weight * kl_loss
grads = tape.gradient(total_loss, self.trainable_weights)
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
# Update metrics
self.reconstruction_loss_tracker.update_state(reconstruction_loss)
self.kl_loss_tracker.update_state(kl_loss)
return {
"loss": total_loss,
"reconstruction_loss": self.reconstruction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result()
}
def test_step(self, data):
z_mean, z_log_var, z = self.encoder(data)
reconstruction = self.decoder(z)
# Flatten the images
data_flat = tf.reshape(data, [tf.shape(data)[0], -1])
reconstruction_flat = tf.reshape(reconstruction, [tf.shape(reconstruction)[0], -1])
# Compute per-sample reconstruction loss
loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=False, reduction=tf.keras.losses.Reduction.NONE)
per_sample_loss = loss_fn(data_flat, reconstruction_flat)
reconstruction_loss = tf.reduce_mean(per_sample_loss)
# KL divergence loss
kl_loss = -0.5 * tf.reduce_mean(
tf.reduce_sum(
1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var),
axis=1
)
)
total_loss = reconstruction_loss + self.kl_weight * kl_loss
# Update metrics
self.reconstruction_loss_tracker.update_state(reconstruction_loss)
self.kl_loss_tracker.update_state(kl_loss)
return {
"loss": total_loss,
"reconstruction_loss": self.reconstruction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result()
}
@property
def metrics(self):
return [self.reconstruction_loss_tracker, self.kl_loss_tracker]
import matplotlib.pyplot as plt
# Funkcja do trenowania VAE
def train_vae(input_shape, latent_dim, x_train, x_val, epochs=30, batch_size=128, dataset_name='Dataset'):
# Build Encoder
encoder, x_shape = build_conv_encoder(input_shape, latent_dim)
# Build Decoder
decoder = build_conv_decoder(x_shape, latent_dim, input_shape)
# Initialize VAE
vae = VAE(encoder, decoder, kl_weight=1.0)
vae.compile(optimizer='adam')
# Train VAE
history = vae.fit(
x_train,
epochs=epochs,
batch_size=batch_size,
validation_data=(x_val, None)
)
# Wizualizacja strat podczas treningu
plt.figure(figsize=(8,6))
plt.plot(history.history['loss'], label='Całkowita strata treningowa')
plt.plot(history.history['reconstruction_loss'], label='Strata rekonstrukcji treningowa')
plt.plot(history.history['kl_loss'], label='Dywergencja KL treningowa')
plt.xlabel('Epoka')
plt.ylabel('Strata')
plt.legend()
plt.title(f'Straty podczas treningu VAE na {dataset_name}')
plt.show()
# Rekonstrukcja obrazów testowych
x_reconstructed = vae.predict(x_val[:10])
# Funkcja do wyświetlania oryginalnych i zrekonstruowanych obrazów
plot_reconstructions(x_val[:10], x_reconstructed, input_shape)
return vae
# Funkcja do wizualizacji rekonstrukcji
def plot_reconstructions(x_original, x_reconstructed, input_shape, num_images=10):
plt.figure(figsize=(20, 4))
for i in range(num_images):
# Wyświetlenie oryginalnych obrazów
ax = plt.subplot(2, num_images, i + 1)
if input_shape[-1] == 1:
plt.imshow(x_original[i].squeeze(), cmap='gray')
else:
plt.imshow(x_original[i])
ax.axis('off')
# Wyświetlenie zrekonstruowanych obrazów
ax = plt.subplot(2, num_images, i + 1 + num_images)
if x_reconstructed.shape[-1] == 1:
plt.imshow(x_reconstructed[i].squeeze(), cmap='gray')
else:
plt.imshow(x_reconstructed[i])
ax.axis('off')
plt.tight_layout()
plt.show()
x_celeba shape: (21000, 64, 64, 3)
x_val_celeba shape: (9000, 64, 64, 3)
# Trenowanie VAE na CelebA
vae_celeba = train_vae(
input_shape=(64, 64, 3),
latent_dim=64,
x_train=x_train_celeba,
x_val=x_val_celeba,
epochs=12,
batch_size=128,
dataset_name='CelebA'
)
# Trening rozpoczyna się poprawnie jednak kończy się złym efektem i nie odtwarza danych, skuteczność wytrenowanych modeli jest równa zeroASSISTANT
It appears you're encountering issues with training your Variational Autoencoder (VAE) on the CelebA dataset, where the model fails to produce meaningful reconstructions. Let's delve into potential reasons for this behavior and explore strategies to improve your model's performance.
### **1. Loss Function Selection**
**Issue:**
Currently, you're using the Binary Crossentropy (BCE) loss for reconstruction:
```python
loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=False, reduction=tf.keras.losses.Reduction.NONE)
```
While BCE works well for binary or grayscale images like MNIST (where pixel values are between 0 and 1), it's less suitable for complex, multi-channel images like those in the CelebA dataset. For RGB images, Mean Squared Error (MSE) or Mean Absolute Error (MAE) often yield better results.
**Solution:**
Switch the reconstruction loss from BCE to MSE. Here's how you can modify both the `train_step` and `test_step` methods in your `VAE` class:
```python
class VAE(Model):
# ... [existing code] ...
def train_step(self, data):
with tf.GradientTape() as tape:
z_mean, z_log_var, z = self.encoder(data)
reconstruction = self.decoder(z)
# Compute per-sample reconstruction loss using MSE
reconstruction_loss_fn = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.NONE)
per_sample_loss = reconstruction_loss_fn(data, reconstruction)
# Sum over the image dimensions
reconstruction_loss = tf.reduce_mean(tf.reduce_sum(per_sample_loss, axis=[1,2,3]))
# KL divergence loss
kl_loss = -0.5 * tf.reduce_mean(
tf.reduce_sum(
1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var),
axis=1
)
)
total_loss = reconstruction_loss + self.kl_weight * kl_loss
grads = tape.gradient(total_loss, self.trainable_weights)
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
# Update metrics
self.reconstruction_loss_tracker.update_state(reconstruction_loss)
self.kl_loss_tracker.update_state(kl_loss)
return {
"loss": total_loss,
"reconstruction_loss": self.reconstruction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result()
}
def test_step(self, data):
z_mean, z_log_var, z = self.encoder(data)
reconstruction = self.decoder(z)
# Compute per-sample reconstruction loss using MSE
reconstruction_loss_fn = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.NONE)
per_sample_loss = reconstruction_loss_fn(data, reconstruction)
# Sum over the image dimensions
reconstruction_loss = tf.reduce_mean(tf.reduce_sum(per_sample_loss, axis=[1,2,3]))
# KL divergence loss
kl_loss = -0.5 * tf.reduce_mean(
tf.reduce_sum(
1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var),
axis=1
)
)
total_loss = reconstruction_loss + self.kl_weight * kl_loss
# Update metrics
self.reconstruction_loss_tracker.update_state(reconstruction_loss)
self.kl_loss_tracker.update_state(kl_loss)
return {
"loss": total_loss,
"reconstruction_loss": self.reconstruction_loss_tracker.result(),
"kl_loss": self.kl_loss_tracker.result()
}
```
**Why This Helps:**
MSE measures the average squared difference between the estimated values and the actual value, which is more appropriate for images with continuous pixel values like those in CelebA. This change should lead to smoother and more accurate reconstructions.
### **2. Model Architecture Enhancements**
**Issue:**
The current encoder and decoder architectures are relatively simple for handling complex facial images in CelebA.
**Solution:**
Enhance the architecture by adding more convolutional layers and increasing the number of filters. This provides the model with greater capacity to capture intricate patterns in the data.
**Revised Encoder:**
```python
def build_conv_encoder(input_shape, latent_dim):
input_img = Input(shape=input_shape)
x = Conv2D(64, kernel_size=3, strides=2, activation='relu', padding='same')(input_img) # 64x64x64
x = Conv2D(128, kernel_size=3, strides=2, activation='relu', padding='same')(x) # 32x32x128
x = Conv2D(256, kernel_size=3, strides=2, activation='relu', padding='same')(x) # 16x16x256
x = Flatten()(x)
x = Dense(512, activation='relu')(x)
z_mean = Dense(latent_dim, name='z_mean')(x)
z_log_var = Dense(latent_dim, name='z_log_var')(x)
z = Sampling()([z_mean, z_log_var])
encoder = Model(input_img, [z_mean, z_log_var, z], name='encoder')
return encoder, (16, 16, 256)
```
**Revised Decoder:**
```python
def build_conv_decoder(x_shape, latent_dim, output_shape):
latent_inputs = Input(shape=(latent_dim,))
x = Dense(np.prod(x_shape), activation='relu')(latent_inputs)
x = Reshape(x_shape)(x) # 16x16x256
x = Conv2DTranspose(256, kernel_size=3, strides=2, activation='relu', padding='same')(x) # 32x32x256
x = Conv2DTranspose(128, kernel_size=3, strides=2, activation='relu', padding='same')(x) # 64x64x128
x = Conv2DTranspose(64, kernel_size=3, strides=1, activation='relu', padding='same')(x) # 64x64x64
outputs = Conv2DTranspose(
output_shape[2],
kernel_size=3,
activation='sigmoid',
padding='same'
)(x) # 64x64x3
decoder = Model(latent_inputs, outputs, name='decoder')
return decoder
```
**Why This Helps:**
By increasing the number of filters and adding more layers, the encoder can extract more detailed features from the images, and the decoder can better reconstruct the high-fidelity images from the latent space.
### **3. Adjusting the KL Divergence Weight**
**Issue:**
The KL divergence term in the loss function encourages the latent space to follow a standard normal distribution. However, if its weight is too high, it can overpower the reconstruction loss, leading the model to ignore the latent codes.
**Solution:**
Consider reducing the `kl_weight` or implementing a **KL annealing** strategy, where the weight starts low and gradually increases during training. This allows the model to first focus on reconstructing the data before regularizing the latent space.
**Example of Reducing `kl_weight`:**
```python
vae = VAE(encoder, decoder, kl_weight=0.1) # Reduced from 1.0 to 0.1
```
**Implementing KL Annealing:**
```python
class VAE(Model):
def __init__(self, encoder, decoder, kl_weight=0.0, kl_annealing_rate=0.001, **kwargs):
super(VAE, self).__init__(**kwargs)
self.encoder = encoder
self.decoder = decoder
self.kl_weight = kl_weight
self.kl_annealing_rate = kl_annealing_rate
# ... [existing code] ...
def train_step(self, data):
# ... [existing code up to loss calculation] ...
# Implement KL annealing
self.kl_weight = min(self.kl_weight + self.kl_annealing_rate, 1.0)
total_loss = reconstruction_loss + self.kl_weight * kl_loss
# ... [rest of the code] ...
```
**Why This Helps:**
Balancing the reconstruction and KL losses ensures that the model learns meaningful latent representations without compromising reconstruction quality.
### **4. Increasing Training Epochs and Monitoring Training**
**Issue:**
Training for only 12 epochs might not be sufficient for the model to converge, especially with a more complex architecture.
**Solution:**
Increase the number of training epochs and monitor the training and validation loss curves to ensure the model is learning appropriately.
```python
vae_celeba = train_vae(
input_shape=(64, 64, 3),
latent_dim=128, # Consider increasing latent dimensions
x_train=x_train_celeba,
x_val=x_val_celeba,
epochs=50, # Increased from 12 to 50
batch_size=128,
dataset_name='CelebA'
)
```
**Why This Helps:**
More epochs allow the model ample time to learn complex representations from the data. Additionally, monitoring loss curves can help identify issues like overfitting or underfitting.
### **5. Enhancing the Latent Space Dimension**
**Issue:**
A latent dimension of 64 might be too restrictive for capturing the variability in CelebA images.
**Solution:**
Increase the latent dimension to allow the model to encode more information.
```python
vae_celeba = train_vae(
input_shape=(64, 64, 3),
latent_dim=128, # Increased from 64 to 128
x_train=x_train_celeba,
x_val=x_val_celeba,
epochs=50,
batch_size=128,
dataset_name='CelebA'
)
```
**Why This Helps:**
A higher latent dimensionality provides the model with more capacity to capture diverse features in the data, potentially leading to better reconstructions.
### **6. Implementing Batch Normalization and Other Enhancements**
**Issue:**
Normalization layers can stabilize and accelerate training by reducing internal covariate shift.
**Solution:**
Incorporate Batch Normalization layers in both the encoder and decoder.
**Modified Encoder with Batch Normalization:**
```python
def build_conv_encoder(input_shape, latent_dim):
input_img = Input(shape=input_shape)
x = Conv2D(64, kernel_size=3, strides=2, padding='same')(input_img)
x = BatchNormalization()(x)
x = Activation('relu')(x)
x = Conv2D(128, kernel_size=3, strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = Activation('relu')(x)
x = Conv2D(256, kernel_size=3, strides=2, padding='same')(x)
x = BatchNormalization()(x)
x = Activation('relu')(x)
x = Flatten()(x)
x = Dense(512)(x)
x = BatchNormalization()(x)
x = Activation('relu')(x)
z_mean = Dense(latent_dim, name='z_mean')(x)
z_log_var = Dense(latent_dim, name='z_log_var')(x)
z = Sampling()([z_mean, z_log_var])
encoder = Model(input_img, [z_mean, z_log_var, z], name='encoder')
return encoder, (16, 16, 256)
```
**Modified Decoder with Batch Normalization:**
```python
def build_conv_decoder(x_shape, latent_dim, output_shape):
latent_inputs = Input(shape=(latent_dim,))
x = Dense(np.prod(x_shape))(latent_inputs)
x = BatchNormalization()(x)
x = Activation('relu')(x)
x = Reshape(x_shape)(x) # 16x16x256
x = Conv2DTranspose(256, kernel_size=3, strides=2, padding='same')(x) # 32x32x256
x = BatchNormalization()(x)
x = Activation('relu')(x)
x = Conv2DTranspose(128, kernel_size=3, strides=2, padding='same')(x) # 64x64x128
x = BatchNormalization()(x)
x = Activation('relu')(x)
x = Conv2DTranspose(64, kernel_size=3, strides=1, padding='same')(x) # 64x64x64
x = BatchNormalization()(x)
x = Activation('relu')(x)
outputs = Conv2DTranspose(
output_shape[2],
kernel_size=3,
activation='sigmoid',
padding='same'
)(x) # 64x64x3
decoder = Model(latent_inputs, outputs, name='decoder')
return decoder
```
**Why This Helps:**
Batch Normalization can improve training stability and convergence speed by normalizing the inputs to each layer, reducing the risk of vanishing or exploding gradients.
### **7. Visualizing the Latent Space**
**Issue:**
Understanding how well the latent space captures the data distribution can provide insights into model performance.
**Solution:**
Use techniques like t-SNE or PCA to visualize the latent vectors. Additionally, generate new images by sampling from the latent space to assess the generative capabilities of the VAE.
**Sample Code for Latent Space Visualization:**
```python
from sklearn.manifold import TSNE
# Get latent vectors
z_mean, _, _ = vae_celeba.encoder.predict(x_val_celeba[:1000])
# Reduce dimensionality to 2D
tsne = TSNE(n_components=2, random_state=42)
z_tsne = tsne.fit_transform(z_mean)
# Plot
plt.figure(figsize=(10, 8))
plt.scatter(z_tsne[:, 0], z_tsne[:, 1], alpha=0.5)
plt.title('t-SNE of VAE Latent Space on CelebA')
plt.xlabel('Dimension 1')
plt.ylabel('Dimension 2')
plt.show()
```
**Why This Helps:**
Visualizing the latent space can reveal clustering patterns, continuity, and how well different classes are separated, providing qualitative insights into the model's representation learning.
### **8. Final Training and Evaluation**
After implementing the above recommendations, retrain your VAE:
```python
vae_celeba = train_vae(
input_shape=(64, 64, 3),
latent_dim=128,
x_train=x_train_celeba,
x_val=x_val_celeba,
epochs=50,
batch_size=128,
dataset_name='CelebA'
)
```
**Key Points to Monitor:**
- **Reconstruction Quality:** Ensure that reconstructed images become progressively clearer and more detailed over epochs.
- **Loss Curves:** Monitor both reconstruction and KL losses. Ideally, reconstruction loss should decrease while KL loss stabilizes without dominating the total loss.
- **Latent Space Distribution:** Check if the latent vectors follow a standard normal distribution as intended.
### **9. Additional Tips**
- **Learning Rate Scheduling:** Implement learning rate decay or use adaptive optimizers like Adam with appropriate learning rates to facilitate better convergence.
- **Regularization:** Experiment with techniques like dropout or L2 regularization to prevent overfitting, especially with increased model complexity.
- **Data Augmentation:** Apply data augmentation techniques (e.g., random flips, rotations) to make the model more robust and prevent overfitting.
### **Conclusion**
By addressing the loss function suitability, enhancing the model architecture, balancing the KL divergence, and implementing effective training strategies, your VAE should achieve improved performance on the CelebA dataset. Remember that training deep generative models can be sensitive to hyperparameters and architecture choices, so iterative experimentation is key to achieving optimal results.
If you continue to experience issues after these adjustments, consider sharing specific error messages or visualization outputs for more targeted assistance.