ganime/model/vqgan_clean/experimental/net2net_v3.py

import numpy as np
import tensorflow as tf
import tensorflow_addons as tfa
from ganime.configs.model_configs import GPTConfig, ModelConfig
from ganime.model.vqgan_clean.experimental.transformer import Transformer
from ganime.model.vqgan_clean.vqgan import VQGAN
from ganime.trainer.warmup.cosine import WarmUpCosine
from tensorflow import keras
from tensorflow.keras import Model, layers
from ganime.model.vqgan_clean.losses.losses import Losses
from ganime.trainer.warmup.base import create_warmup_scheduler
from ganime.visualization.images import unnormalize_if_necessary


class Net2Net(Model):
    def __init__(
        self,
        transformer_config: GPTConfig,
        first_stage_config: ModelConfig,
        trainer_config,
        num_replicas: int = 1,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.first_stage_model = VQGAN(**first_stage_config)
        self.transformer = Transformer(transformer_config)

        if "checkpoint_path" in transformer_config:
            print(f"Restoring weights from {transformer_config['checkpoint_path']}")
            self.load_weights(transformer_config["checkpoint_path"])

        self.scheduled_lrs = create_warmup_scheduler(
            trainer_config, num_devices=num_replicas
        )

        optimizer = tfa.optimizers.AdamW(
            learning_rate=self.scheduled_lrs, weight_decay=1e-4
        )
        self.compile(
            optimizer=optimizer,
            # loss=self.loss_fn,
            # run_eagerly=True,
        )

        self.n_frames_before = trainer_config["n_frames_before"]

        # Gradient accumulation
        self.gradient_accumulation = [
            tf.Variable(tf.zeros_like(v, dtype=tf.float32), trainable=False)
            for v in self.transformer.trainable_variables
        ]
        self.accumulation_size = trainer_config["accumulation_size"]

        # Losses
        self.perceptual_loss_weight = trainer_config["perceptual_loss_weight"]
        losses = Losses(num_replicas=num_replicas)
        self.scce_loss = losses.scce_loss
        self.perceptual_loss = losses.perceptual_loss

        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.scce_loss_tracker = keras.metrics.Mean(name="scce_loss")
        self.perceptual_loss_tracker = keras.metrics.Mean(name="perceptual_loss")

        self.epoch = 0
        self.stop_ground_truth_after_epoch = trainer_config[
            "stop_ground_truth_after_epoch"
        ]

    def apply_accu_gradients(self):
        # apply accumulated gradients
        self.optimizer.apply_gradients(
            zip(self.gradient_accumulation, self.transformer.trainable_variables)
        )

        # reset
        for i in range(len(self.gradient_accumulation)):
            self.gradient_accumulation[i].assign(
                tf.zeros_like(self.transformer.trainable_variables[i], dtype=tf.float32)
            )

    @property
    def metrics(self):
        # We list our `Metric` objects here so that `reset_states()` can be
        # called automatically at the start of each epoch
        # or at the start of `evaluate()`.
        # If you don't implement this property, you have to call
        # `reset_states()` yourself at the time of your choosing.
        return [
            self.total_loss_tracker,
            self.scce_loss_tracker,
            self.perceptual_loss_tracker,
        ]

    @tf.function()
    def encode_to_z(self, x):
        quant_z, indices, quantized_loss = self.first_stage_model.encode(x)

        batch_size = tf.shape(quant_z)[0]

        indices = tf.reshape(indices, shape=(batch_size, -1))
        return quant_z, indices

    def call(self, inputs, training=False, mask=None, return_losses=False):

        return self.predict_video(inputs, training, return_losses)

    def predict(self, data, sample=False, temperature=1.0):
        video = self.predict_video(
            data,
            training=False,
            return_losses=False,
            sample=sample,
            temperature=temperature,
        )
        video = unnormalize_if_necessary(video)
        return video

    def get_remaining_frames(self, inputs):
        if "remaining_frames" in inputs:
            remaining_frames = inputs["remaining_frames"]
        else:
            raise NotImplementedError
        remaining_frames = tf.cast(remaining_frames, tf.int64)
        return remaining_frames

    # @tf.function()
    def predict_video(
        self, inputs, training=False, return_losses=False, sample=False, temperature=1.0
    ):
        first_frame = inputs["first_frame"]
        last_frame = inputs["last_frame"]
        n_frames = tf.reduce_max(inputs["n_frames"])
        remaining_frames = self.get_remaining_frames(inputs)

        try:
            ground_truth = inputs["y"]
        except AttributeError:
            ground_truth = None

        previous_frames = tf.expand_dims(first_frame, axis=1)

        predictions = tf.TensorArray(
            tf.float32, size=0, dynamic_size=True, clear_after_read=False
        )

        quant_last, indices_last = self.encode_to_z(last_frame)

        total_loss = tf.constant(0.0)
        scce_loss = tf.constant(0.0)
        perceptual_loss = tf.constant(0.0)

        current_frame_index = tf.constant(1)
        while tf.less(current_frame_index, n_frames):
            tf.autograph.experimental.set_loop_options(
                shape_invariants=[
                    (previous_frames, tf.TensorShape([None, None, None, None, 3]))
                ],
            )

            if ground_truth is not None:
                target_frame = ground_truth[:, current_frame_index]
            else:
                target_frame = None

            y_pred, losses = self.predict_next_frame(
                remaining_frames[:, current_frame_index],
                previous_frames,
                last_frame,
                indices_last,
                quant_last,
                target_frame=target_frame,
                training=training,
                sample=sample,
                temperature=temperature,
            )
            predictions = predictions.write(current_frame_index, y_pred)

            if training and self.epoch < self.stop_ground_truth_after_epoch:
                start_index = tf.math.maximum(
                    0, current_frame_index - self.n_frames_before
                )
                previous_frames = ground_truth[
                    :, start_index + 1 : current_frame_index + 1
                ]
            else:
                previous_frames = predictions.stack()
                previous_frames = tf.transpose(previous_frames, (1, 0, 2, 3, 4))
                previous_frames = previous_frames[:, -self.n_frames_before :]

            current_frame_index = tf.add(current_frame_index, 1)
            total_loss = tf.add(total_loss, losses[0])
            scce_loss = tf.add(scce_loss, losses[1])
            perceptual_loss = tf.add(perceptual_loss, losses[2])

        predictions = predictions.stack()
        predictions = tf.transpose(predictions, (1, 0, 2, 3, 4))

        total_loss = tf.divide(total_loss, tf.cast(n_frames, tf.float32))
        scce_loss = tf.divide(scce_loss, tf.cast(n_frames, tf.float32))
        perceptual_loss = tf.divide(perceptual_loss, tf.cast(n_frames, tf.float32))

        if return_losses:
            return predictions, total_loss, scce_loss, perceptual_loss
        else:
            return predictions

    def predict_next_frame(
        self,
        remaining_frames,
        previous_frames,
        last_frame,
        indices_last,
        quant_last,
        target_frame=None,
        training=False,
        sample=False,
        temperature=1.0,
    ):
        # previous frames is of shape (batch_size, n_frames, height, width, 3)
        previous_frames = tf.transpose(previous_frames, (1, 0, 2, 3, 4))
        # previous frames is now of shape (n_frames, batch_size, height, width, 3)

        indices_previous = tf.map_fn(
            lambda x: self.encode_to_z(x)[1],
            previous_frames,
            fn_output_signature=tf.int64,
        )

        # indices is of shape (n_frames, batch_size, n_z)
        indices_previous = tf.transpose(indices_previous, (1, 0, 2))
        # indices is now of shape (batch_size, n_frames, n_z)
        batch_size, n_frames, n_z = (
            tf.shape(indices_previous)[0],
            tf.shape(indices_previous)[1],
            tf.shape(indices_previous)[2],
        )
        indices_previous = tf.reshape(
            indices_previous, shape=(batch_size, n_frames * n_z)
        )

        if target_frame is not None:
            _, target_indices = self.encode_to_z(target_frame)
        else:
            target_indices = None

        if training:
            next_frame, losses = self.train_predict_next_frame(
                remaining_frames,
                indices_last,
                indices_previous,
                target_indices=target_indices,
                target_frame=target_frame,
                quant_shape=tf.shape(quant_last),
                indices_shape=tf.shape(indices_last),
            )
        else:
            next_frame, losses = self.predict_next_frame_body(
                remaining_frames,
                indices_last,
                indices_previous,
                target_indices=target_indices,
                target_frame=target_frame,
                quant_shape=tf.shape(quant_last),
                indices_shape=tf.shape(indices_last),
                sample=sample,
                temperature=temperature,
            )

        return next_frame, losses

    def predict_next_frame_body(
        self,
        remaining_frames,
        last_frame_indices,
        previous_frame_indices,
        quant_shape,
        indices_shape,
        target_indices=None,
        target_frame=None,
        sample=False,
        temperature=1.0,
    ):
        logits = self.transformer(
            (remaining_frames, last_frame_indices, previous_frame_indices)
        )
        next_frame = self.convert_logits_to_image(
            logits,
            quant_shape=quant_shape,
            indices_shape=indices_shape,
            sample=sample,
            temperature=temperature,
        )
        if target_indices is not None:
            scce_loss = self.scce_loss(target_indices, logits)
        else:
            scce_loss = 0.0

        if target_frame is not None:
            perceptual_loss = 1.0 * self.perceptual_loss(target_frame, next_frame)
        else:
            perceptual_loss = 0.0

        frame_loss = scce_loss + perceptual_loss

        # self.total_loss_tracker.update_state(frame_loss)
        # self.scce_loss_tracker.update_state(scce_loss)
        # self.perceptual_loss_tracker.update_state(perceptual_loss)

        return next_frame, (frame_loss, scce_loss, perceptual_loss)

    def train_predict_next_frame(
        self,
        remaining_frames,
        last_frame_indices,
        previous_frame_indices,
        quant_shape,
        indices_shape,
        target_indices,
        target_frame,
    ):
        with tf.GradientTape() as tape:
            next_frame, losses = self.predict_next_frame_body(
                remaining_frames=remaining_frames,
                last_frame_indices=last_frame_indices,
                previous_frame_indices=previous_frame_indices,
                target_indices=target_indices,
                quant_shape=quant_shape,
                indices_shape=indices_shape,
                target_frame=target_frame,
                sample=False,
            )
        frame_loss = losses[0]
        # Calculate batch gradients
        gradients = tape.gradient(frame_loss, self.transformer.trainable_variables)

        # Accumulate batch gradients
        for i in range(len(self.gradient_accumulation)):
            self.gradient_accumulation[i].assign_add(tf.cast(gradients[i], tf.float32))

        return next_frame, losses

    def convert_logits_to_image(
        self, logits, quant_shape, indices_shape, sample=False, temperature=1.0
    ):
        if sample:
            array = []
            for i in range(logits.shape[1]):
                sub_logits = logits[:, i]
                sub_logits = sub_logits / temperature
                # sub_logits, _ = tf.math.top_k(sub_logits, k=1)
                probs = tf.keras.activations.softmax(sub_logits)
                probs, probs_index = tf.math.top_k(probs, k=50)
                selection_index = tf.random.categorical(
                    tf.math.log(probs), num_samples=1
                )
                ix = tf.gather_nd(probs_index, selection_index, batch_dims=1)
                ix = tf.reshape(ix, (-1, 1))
                array.append(ix)
            generated_indices = tf.concat(array, axis=-1)
        else:
            probs = tf.keras.activations.softmax(logits)
            _, generated_indices = tf.math.top_k(probs)

        generated_indices = tf.reshape(
            generated_indices,
            indices_shape,
        )
        quant = self.first_stage_model.quantize.get_codebook_entry(
            generated_indices, shape=quant_shape
        )

        return self.first_stage_model.decode(quant)

    def train_step(self, data):

        batch_total_loss, batch_scce_loss, batch_perceptual_loss = 0.0, 0.0, 0.0
        for i in range(self.accumulation_size):
            sub_data = {
                key: value[
                    self.accumulation_size * i : self.accumulation_size * (i + 1)
                ]
                for key, value in data.items()
            }
            _, total_loss, scce_loss, perceptual_loss = self(
                sub_data, training=True, return_losses=True
            )
            batch_total_loss += total_loss
            batch_scce_loss += scce_loss
            batch_perceptual_loss += perceptual_loss

        self.apply_accu_gradients()
        self.total_loss_tracker.update_state(batch_total_loss)
        self.scce_loss_tracker.update_state(batch_scce_loss)
        self.perceptual_loss_tracker.update_state(batch_perceptual_loss)
        self.epoch += 1
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        _, total_loss, scce_loss, perceptual_loss = self(
            data, training=False, return_losses=True
        )

        self.total_loss_tracker.update_state(total_loss)
        self.scce_loss_tracker.update_state(scce_loss)
        self.perceptual_loss_tracker.update_state(perceptual_loss)
        return {m.name: m.result() for m in self.metrics}