<reponame>chris4540/DD2412-adl-proj<gh_stars>0
import numpy as np
import tensorflow as tf
from tensorflow.keras.losses import KLDivergence
from tensorflow.keras.utils import normalize

def kldiv_loss_fn(p_true, p_pred):
    """
    Kullback Leibler divergence loss

    Args:
        p_true:
        p_pred:

    Ref:
    https://www.tensorflow.org/versions/r1.14/api_docs/python/tf/keras/losses/KLDivergence
    """
    return KLDivergence()(p_true, p_pred)

def generator_loss_fn(t_logits, s_logits, temp=1):

	loss = kldiv_loss_fn(
            tf.math.softmax(t_logits / temp) ,
            tf.math.softmax(s_logits / temp))

	g_loss = -loss

	return g_loss

def cross_entropy(logits, onehot_label):
    cce = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
    ret = cce(onehot_label, logits)
    return ret

def knowledge_distil_loss_fn(**kwargs):
    alpha = kwargs['alpha']
    beta = kwargs['beta']
    onehot_label = kwargs['onehot_label']
    t_logits = kwargs['t_logits']
    t_acts = kwargs['t_acts']
    s_logits = kwargs['s_logits']
    s_acts = kwargs['s_acts']
    temp = kwargs.get('temp', 1)
    # ====================================================================

    # 1. cross entopy loss
    if alpha < 1.0:
        ce_loss = cross_entropy(s_logits, onehot_label)
    else:
        ce_loss = 0

    # 2. KL Divergence
    if alpha > 0.0:
        kl_div_loss = kldiv_loss_fn(
                tf.math.softmax(t_logits / temp) ,
                tf.math.softmax(s_logits / temp))
    else:
        kl_div_loss = 0.0

    # attention loss
    att_loss = 0.0
    if beta != 0.0:
        for t_act, s_act in zip(t_acts, s_acts):
            att_loss += attention_loss(t_act, s_act)

    ret = (1 - alpha) * ce_loss + (2*alpha*temp*temp) * kl_div_loss + beta* att_loss
    return ret

def student_loss_fn(t_logits, t_acts, s_logits, s_acts, beta, temp=1):
    """
    The student loss function used in
        - zero-shot learning

    See Section 3.2, Eq. 1

    Args:
        t_logits: Teacher logits
        t_acts:  list of teacher activation layer output
        s_logits: Student logits
        s_acts: list of student activation layer output
        beta: hyper-parameter for tuning the weight of the attention term

    Return:
        loss
    """
    loss = kldiv_loss_fn(
            tf.math.softmax(t_logits / temp) ,
            tf.math.softmax(s_logits / temp))

    if beta != 0.0:
        att_loss = 0.0
        for t_act, s_act in zip(t_acts, s_acts):
            att_loss += attention_loss(t_act, s_act)

        loss += beta * att_loss

    return loss

def __spatial_attention_map(act_tensor, p=2):
    """
    Spatial attention mapping function to map the activation tensor with shape
    (H, W, C) to (H, W).

    We employed:
        sum of absolute values raised to the power of 2

    The f(A_{l}) is the paper of replication

    Args:
        act_tensor: activation tensor with shape (H, W, C)
    Return:
        a spatial attention map with shape (H, W)

    Migration:
        corr. to f_act
    """

    out = tf.pow(act_tensor, p)
    out = tf.reduce_mean(out, axis=-1)
    # flatten it
    out = tf.reshape(out, [out.shape[0], -1])

    # renormalize them
    out = tf.linalg.l2_normalize(out)
    return out

def attention_loss(act1, act2):
    """
    Return the activation loss. The loss is the L2 distances between two
    activation map

    Args:
        act_map_1:
        act_map_2:

    Return:
        a floating point number representing the loss. As we use tensorflow,
        the floating point number would be a number hold in tf.Tensor

    Bug:
        1. use this with beta = 250 will blow up the err and the grad will explode
        2. their implementation not using beta = 250
    Ref:
    https://github.com/szagoruyko/attention-transfer/blob/893df5488f93691799f082a70e2521a9dc2ddf2d/utils.py#L22
    """
    # get the activation map first
    act_map_1 = __spatial_attention_map(act1)
    act_map_2 = __spatial_attention_map(act2)

    if False: # paper impl.
        # calculate vector norm of vectorized matrix
        out = tf.pow(act_map_1 - act_map_2, 2)
        out = tf.reduce_sum(out)
        ret = tf.sqrt(out)
    else:
        # their code impl
        out = tf.pow(act_map_1 - act_map_2, 2)
        out = tf.reduce_mean(out)
        ret = out
    return ret

import numpy as np
import tensorflow as tf
from tensorflow.keras.losses import KLDivergence
from tensorflow.keras.utils import normalize

def kldiv_loss_fn(p_true, p_pred):
    """
    Kullback Leibler divergence loss

    Args:
        p_true:
        p_pred:

    Ref:
    https://www.tensorflow.org/versions/r1.14/api_docs/python/tf/keras/losses/KLDivergence
    """
    return KLDivergence()(p_true, p_pred)

def generator_loss_fn(t_logits, s_logits, temp=1):

	loss = kldiv_loss_fn(
            tf.math.softmax(t_logits / temp) ,
            tf.math.softmax(s_logits / temp))

	g_loss = -loss

	return g_loss

def cross_entropy(logits, onehot_label):
    cce = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
    ret = cce(onehot_label, logits)
    return ret

def knowledge_distil_loss_fn(**kwargs):
    alpha = kwargs['alpha']
    beta = kwargs['beta']
    onehot_label = kwargs['onehot_label']
    t_logits = kwargs['t_logits']
    t_acts = kwargs['t_acts']
    s_logits = kwargs['s_logits']
    s_acts = kwargs['s_acts']
    temp = kwargs.get('temp', 1)
    # ====================================================================

    # 1. cross entopy loss
    if alpha < 1.0:
        ce_loss = cross_entropy(s_logits, onehot_label)
    else:
        ce_loss = 0

    # 2. KL Divergence
    if alpha > 0.0:
        kl_div_loss = kldiv_loss_fn(
                tf.math.softmax(t_logits / temp) ,
                tf.math.softmax(s_logits / temp))
    else:
        kl_div_loss = 0.0

    # attention loss
    att_loss = 0.0
    if beta != 0.0:
        for t_act, s_act in zip(t_acts, s_acts):
            att_loss += attention_loss(t_act, s_act)

    ret = (1 - alpha) * ce_loss + (2*alpha*temp*temp) * kl_div_loss + beta* att_loss
    return ret

def student_loss_fn(t_logits, t_acts, s_logits, s_acts, beta, temp=1):
    """
    The student loss function used in
        - zero-shot learning

    See Section 3.2, Eq. 1

    Args:
        t_logits: Teacher logits
        t_acts:  list of teacher activation layer output
        s_logits: Student logits
        s_acts: list of student activation layer output
        beta: hyper-parameter for tuning the weight of the attention term

    Return:
        loss
    """
    loss = kldiv_loss_fn(
            tf.math.softmax(t_logits / temp) ,
            tf.math.softmax(s_logits / temp))

    if beta != 0.0:
        att_loss = 0.0
        for t_act, s_act in zip(t_acts, s_acts):
            att_loss += attention_loss(t_act, s_act)

        loss += beta * att_loss

    return loss

def __spatial_attention_map(act_tensor, p=2):
    """
    Spatial attention mapping function to map the activation tensor with shape
    (H, W, C) to (H, W).

    We employed:
        sum of absolute values raised to the power of 2

    The f(A_{l}) is the paper of replication

    Args:
        act_tensor: activation tensor with shape (H, W, C)
    Return:
        a spatial attention map with shape (H, W)

    Migration:
        corr. to f_act
    """

    out = tf.pow(act_tensor, p)
    out = tf.reduce_mean(out, axis=-1)
    # flatten it
    out = tf.reshape(out, [out.shape[0], -1])

    # renormalize them
    out = tf.linalg.l2_normalize(out)
    return out

def attention_loss(act1, act2):
    """
    Return the activation loss. The loss is the L2 distances between two
    activation map

    Args:
        act_map_1:
        act_map_2:

    Return:
        a floating point number representing the loss. As we use tensorflow,
        the floating point number would be a number hold in tf.Tensor

    Bug:
        1. use this with beta = 250 will blow up the err and the grad will explode
        2. their implementation not using beta = 250
    Ref:
    https://github.com/szagoruyko/attention-transfer/blob/893df5488f93691799f082a70e2521a9dc2ddf2d/utils.py#L22
    """
    # get the activation map first
    act_map_1 = __spatial_attention_map(act1)
    act_map_2 = __spatial_attention_map(act2)

    if False: # paper impl.
        # calculate vector norm of vectorized matrix
        out = tf.pow(act_map_1 - act_map_2, 2)
        out = tf.reduce_sum(out)
        ret = tf.sqrt(out)
    else:
        # their code impl
        out = tf.pow(act_map_1 - act_map_2, 2)
        out = tf.reduce_mean(out)
        ret = out
    return ret
Kullback Leibler divergence loss      Args:         p_true:         p_pred:      Ref:     https://www.tensorflow.org/versions/r1.14/api_docs/python/tf/keras/losses/KLDivergence # ==================================================================== # 1. cross entopy loss # 2. KL Divergence # attention loss The student loss function used in         - zero-shot learning      See Section 3.2, Eq. 1      Args:         t_logits: Teacher logits         t_acts:  list of teacher activation layer output         s_logits: Student logits         s_acts: list of student activation layer output         beta: hyper-parameter for tuning the weight of the attention term      Return:         loss Spatial attention mapping function to map the activation tensor with shape     (H, W, C) to (H, W).      We employed:         sum of absolute values raised to the power of 2      The f(A_{l}) is the paper of replication      Args:         act_tensor: activation tensor with shape (H, W, C)     Return:         a spatial attention map with shape (H, W)      Migration:         corr. to f_act # flatten it # renormalize them Return the activation loss. The loss is the L2 distances between two     activation map      Args:         act_map_1:         act_map_2:      Return:         a floating point number representing the loss. As we use tensorflow,         the floating point number would be a number hold in tf.Tensor      Bug:         1. use this with beta = 250 will blow up the err and the grad will explode         2. their implementation not using beta = 250     Ref:     https://github.com/szagoruyko/attention-transfer/blob/893df5488f93691799f082a70e2521a9dc2ddf2d/utils.py#L22 # get the activation map first # paper impl. # calculate vector norm of vectorized matrix # their code impl