Lab4.py - master - Durbolay/Lab4

Lab4.py
136 строк6 KB
Durbolay
create Lab4.py
19 ноя 2025, 16:18
19 ноя 2025, 16:1822c7817
# Исправленная ЛР4: VAE для Fashion-MNIST (готово для Google Colab)
import tensorflow as tf, numpy as np, matplotlib.pyplot as plt
from tensorflow.keras import layers, models

# Загрузка и подготовка данных
(x_train,_),(x_test,_)=tf.keras.datasets.fashion_mnist.load_data()
x_train = x_train.astype("float32") / 255.0
x_test  = x_test.astype("float32")  / 255.0
x_train = np.expand_dims(x_train, -1)
x_test  = np.expand_dims(x_test, -1)

latent_dim = 2  # удобно для визуализации

# Sampling layer (reparameterization trick)
class Sampling(layers.Layer):
    def call(self, inputs):
        z_mean, z_log_var = inputs
        eps = tf.random.normal(shape=tf.shape(z_mean))
        return z_mean + tf.exp(0.5 * z_log_var) * eps

# Энкодер
enc_in = layers.Input(shape=(28,28,1))
x = layers.Conv2D(32,3,activation='relu',strides=2,padding='same')(enc_in)
x = layers.Conv2D(64,3,activation='relu',strides=2,padding='same')(x)
x = layers.Flatten()(x)
x = layers.Dense(16, activation='relu')(x)
z_mean = layers.Dense(latent_dim, name='z_mean')(x)
z_log  = layers.Dense(latent_dim, name='z_log_var')(x)
z = Sampling()([z_mean, z_log])
encoder = models.Model(enc_in, [z_mean, z_log, z], name='encoder')
encoder.summary()

# Декодер
lat_in = layers.Input(shape=(latent_dim,))
x = layers.Dense(7*7*64, activation='relu')(lat_in)
x = layers.Reshape((7,7,64))(x)
x = layers.Conv2DTranspose(64,3,activation='relu',strides=2,padding='same')(x)
x = layers.Conv2DTranspose(32,3,activation='relu',strides=2,padding='same')(x)
dec_out = layers.Conv2DTranspose(1,3,activation='sigmoid',padding='same')(x)
decoder = models.Model(lat_in, dec_out, name='decoder')
decoder.summary()

# Кастомный VAE с train_step и call (call нужен для валидации/вызова model(x))
class VAE(tf.keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        # метрики (опционально)
        self.total_loss_tracker = tf.keras.metrics.Mean(name="loss")
        self.recon_loss_tracker = tf.keras.metrics.Mean(name="recon_loss")
        self.kl_loss_tracker = tf.keras.metrics.Mean(name="kl_loss")

    @property
    def metrics(self):
        return [self.total_loss_tracker, self.recon_loss_tracker, self.kl_loss_tracker]

    # call позволяет model(x) возвращать реконструкцию (нужно для evaluate/validation)
    def call(self, inputs, training=False):
        z_mean, z_log_var, z = self.encoder(inputs, training=training)
        return self.decoder(z, training=training)

    def train_step(self, data):
        if isinstance(data, tuple):
            x = data[0]
        else:
            x = data
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(x, training=True)
            reconstruction = self.decoder(z, training=True)
            # reconstruction loss (binary crossentropy per pixel)
            recon_loss = tf.reduce_mean(
                tf.keras.losses.binary_crossentropy(x, reconstruction)
            ) * 28 * 28
            # KL divergence
            kl_loss = -0.5 * tf.reduce_mean(1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            total_loss = recon_loss + kl_loss
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))

        # Обновление метрик
        self.total_loss_tracker.update_state(total_loss)
        self.recon_loss_tracker.update_state(recon_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {"loss": self.total_loss_tracker.result(),
                "reconstruction_loss": self.recon_loss_tracker.result(),
                "kl_loss": self.kl_loss_tracker.result()}

    # тестовый шаг (чтобы корректно считать валидацию)
    def test_step(self, data):
        if isinstance(data, tuple):
            x = data[0]
        else:
            x = data
        z_mean, z_log_var, z = self.encoder(x, training=False)
        reconstruction = self.decoder(z, training=False)
        recon_loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(x, reconstruction)) * 28 * 28
        kl_loss = -0.5 * tf.reduce_mean(1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
        total_loss = recon_loss + kl_loss

        # Обновление метрик (валидация)
        self.total_loss_tracker.update_state(total_loss)
        self.recon_loss_tracker.update_state(recon_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {"loss": self.total_loss_tracker.result(),
                "reconstruction_loss": self.recon_loss_tracker.result(),
                "kl_loss": self.kl_loss_tracker.result()}

# Создаём и тренируем VAE
vae = VAE(encoder, decoder)
vae.compile(optimizer=tf.keras.optimizers.Adam())

# ВАЖНО: корректное validation_data — подаём x_test как и для обучения (не (x_test, None))
history = vae.fit(x_train, epochs=15, batch_size=128, validation_data=(x_test,))

# Визуализация латентного пространства (используя z_mean)
z_mean, z_log, z_samples = encoder.predict(x_test, batch_size=256)
plt.figure(figsize=(6,6))
plt.scatter(z_mean[:,0], z_mean[:,1], s=1)
plt.title('Latent space (z_mean)')
plt.show()

# Генерация сетки изображений из латентного пространства (для latent_dim=2)
grid_x = np.linspace(-3, 3, 15)
grid_y = np.linspace(-3, 3, 15)[::-1]
figure = np.zeros((28 * len(grid_y), 28 * len(grid_x)))
for i, yi in enumerate(grid_y):
    for j, xi in enumerate(grid_x):
        z_sample = np.array([[xi, yi]])
        x_decoded = decoder.predict(z_sample)
        digit = x_decoded[0].squeeze()
        figure[i * 28: (i + 1) * 28, j * 28: (j + 1) * 28] = digit
plt.figure(figsize=(8,8))
plt.imshow(figure, cmap='gray')
plt.axis('off')
plt.show()