<reponame>Shi-Yupeng/hessian_penalty
import torch



def jaocbian_panelty(G, z, c, w, Q, epsilon=0.1, **kwargs):
    G_z = G(z, c, w=w, Q=Q, multi=True, **kwargs)
    xs_idx1 = torch.max(torch.rand_like(w), dim=1)[1]
    xs_delta = torch.max(torch.rand_like(w), dim=1)[1]
    xs_delta = torch.clamp(xs_delta, 1, w.size(1))
    xs_idx2 = (xs_idx1 + xs_delta) % w.size(1)

    xs_idx1 = xs_idx1.unsqueeze(0).unsqueeze(2)
    xs1 = torch.zeros(1, w.size(0), w.size(1)).to(w.device)
    xs1.scatter_(2, xs_idx1.long(), 1)

    xs_idx2 = xs_idx2.unsqueeze(0).unsqueeze(2)
    xs2 = torch.zeros(1, w.size(0), w.size(1)).to(w.device)
    xs2.scatter_(2, xs_idx2.long(), 1)

    xs1 = xs1 * epsilon
    xs2 = xs2 * epsilon

    xs = torch.cat((xs1, xs2), 0)

    central_first_order = multi_layer_first_directional_derivative(G, z, xs, G_z, c, w, Q, epsilon, **kwargs)
    loss = multi_stack_var_and_reduce(central_first_order)
    return loss

def multi_layer_first_directional_derivative(G, z, xs, G_z, c, w, Q, epsilon, **kwargs):
    G_i = G(z, c, w=w+xs[0], Q=Q, multi=True, **kwargs)
    G_j = G(z, c, w=w+xs[1], Q=Q, multi=True, **kwargs)
    sdd = [[[(Gi - Gz) / epsilon, (Gj - Gz) / epsilon], []] for Gi, Gz, Gj in zip(G_i, G_z, G_j)]
    return sdd

def multi_stack_var_and_reduce(sdds):
    """Iterate over all activations to be regularized, then apply Equation (5) to each."""
    sum_of_penalties = 0
    # weights = [4, 2, 1, 0.5, 1]
    # weights = [0.25, 0.5, 1, 1.5, 1]
    # weights = [1, 1, 1, 1, 1, 1]
    for activ_n in sdds:
        Gi = activ_n[0][0]
        Gj = activ_n[0][1]
        penalty = torch.nn.functional.cosine_similarity(Gi.view(Gi.size(0), -1), Gj.view(Gj.size(0), -1), dim=1)
        # penalty = torch.bmm(Gi.view(Gi.size(0), 1, -1), Gj.view(Gj.size(0), -1, 1))
        penalty = torch.mean(torch.abs(penalty))
        sum_of_penalties += penalty
    return sum_of_penalties
import torch



def jaocbian_panelty(G, z, c, w, Q, epsilon=0.1, **kwargs):
    G_z = G(z, c, w=w, Q=Q, multi=True, **kwargs)
    xs_idx1 = torch.max(torch.rand_like(w), dim=1)[1]
    xs_delta = torch.max(torch.rand_like(w), dim=1)[1]
    xs_delta = torch.clamp(xs_delta, 1, w.size(1))
    xs_idx2 = (xs_idx1 + xs_delta) % w.size(1)

    xs_idx1 = xs_idx1.unsqueeze(0).unsqueeze(2)
    xs1 = torch.zeros(1, w.size(0), w.size(1)).to(w.device)
    xs1.scatter_(2, xs_idx1.long(), 1)

    xs_idx2 = xs_idx2.unsqueeze(0).unsqueeze(2)
    xs2 = torch.zeros(1, w.size(0), w.size(1)).to(w.device)
    xs2.scatter_(2, xs_idx2.long(), 1)

    xs1 = xs1 * epsilon
    xs2 = xs2 * epsilon

    xs = torch.cat((xs1, xs2), 0)

    central_first_order = multi_layer_first_directional_derivative(G, z, xs, G_z, c, w, Q, epsilon, **kwargs)
    loss = multi_stack_var_and_reduce(central_first_order)
    return loss

def multi_layer_first_directional_derivative(G, z, xs, G_z, c, w, Q, epsilon, **kwargs):
    G_i = G(z, c, w=w+xs[0], Q=Q, multi=True, **kwargs)
    G_j = G(z, c, w=w+xs[1], Q=Q, multi=True, **kwargs)
    sdd = [[[(Gi - Gz) / epsilon, (Gj - Gz) / epsilon], []] for Gi, Gz, Gj in zip(G_i, G_z, G_j)]
    return sdd

def multi_stack_var_and_reduce(sdds):
    """Iterate over all activations to be regularized, then apply Equation (5) to each."""
    sum_of_penalties = 0
    # weights = [4, 2, 1, 0.5, 1]
    # weights = [0.25, 0.5, 1, 1.5, 1]
    # weights = [1, 1, 1, 1, 1, 1]
    for activ_n in sdds:
        Gi = activ_n[0][0]
        Gj = activ_n[0][1]
        penalty = torch.nn.functional.cosine_similarity(Gi.view(Gi.size(0), -1), Gj.view(Gj.size(0), -1), dim=1)
        # penalty = torch.bmm(Gi.view(Gi.size(0), 1, -1), Gj.view(Gj.size(0), -1, 1))
        penalty = torch.mean(torch.abs(penalty))
        sum_of_penalties += penalty
    return sum_of_penalties
Iterate over all activations to be regularized, then apply Equation (5) to each. # weights = [4, 2, 1, 0.5, 1] # weights = [0.25, 0.5, 1, 1.5, 1] # weights = [1, 1, 1, 1, 1, 1] # penalty = torch.bmm(Gi.view(Gi.size(0), 1, -1), Gj.view(Gj.size(0), -1, 1))