code/helpers/imle_helpers.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
import imageio
from visual.utils import get_sample_for_visualization, generate_for_NN, generate_images_initial
from torch.utils.data import DataLoader, TensorDataset
from helpers.utils import ZippedDataset, get_cpu_stats_over_ranks


@torch.jit.script
def gaussian_analytical_kl(mu1, mu2, logsigma1, logsigma2):
    return -0.5 + logsigma2 - logsigma1 + 0.5 * (logsigma1.exp() ** 2 + (mu1 - mu2) ** 2) / (logsigma2.exp() ** 2)


@torch.jit.script
def draw_gaussian_diag_samples(mu, logsigma, eps):
    return torch.exp(logsigma) * eps + mu


def get_conv(in_dim, out_dim, kernel_size, stride, padding, zero_bias=True, zero_weights=False, groups=1, scaled=False):
    c = nn.Conv2d(in_dim, out_dim, kernel_size, stride, padding, groups=groups)
    if zero_bias:
        c.bias.data *= 0.0
    if zero_weights:
        c.weight.data *= 0.0
    return c


def get_3x3(in_dim, out_dim, zero_bias=True, zero_weights=False, groups=1, scaled=False):
    return get_conv(in_dim, out_dim, 3, 1, 1, zero_bias, zero_weights, groups=groups, scaled=scaled)


def get_1x1(in_dim, out_dim, zero_bias=True, zero_weights=False, groups=1, scaled=False):
    return get_conv(in_dim, out_dim, 1, 1, 0, zero_bias, zero_weights, groups=groups, scaled=scaled)


def log_prob_from_logits(x):
    """ numerically stable log_softmax implementation that prevents overflow """
    axis = len(x.shape) - 1
    m = x.max(dim=axis, keepdim=True)[0]
    return x - m - torch.log(torch.exp(x - m).sum(dim=axis, keepdim=True))


def const_max(t, constant):
    other = torch.ones_like(t) * constant
    return torch.max(t, other)


def const_min(t, constant):
    other = torch.ones_like(t) * constant
    return torch.min(t, other)


def discretized_mix_logistic_loss(x, l, low_bit=False):
    """ log-likelihood for mixture of discretized logistics, assumes the data has been rescaled to [-1,1] interval """
    # Adapted from https://github.com/openai/pixel-cnn/blob/master/pixel_cnn_pp/nn.py
    xs = [s for s in x.shape]  # true image (i.e. labels) to regress to, e.g. (B,32,32,3)
    ls = [s for s in l.shape]  # predicted distribution, e.g. (B,32,32,100)
    nr_mix = int(ls[-1] / 10)  # here and below: unpacking the params of the mixture of logistics
    logit_probs = l[:, :, :, :nr_mix]
    l = torch.reshape(l[:, :, :, nr_mix:], xs + [nr_mix * 3])
    means = l[:, :, :, :, :nr_mix]
    log_scales = const_max(l[:, :, :, :, nr_mix:2 * nr_mix], -7.)
    coeffs = torch.tanh(l[:, :, :, :, 2 * nr_mix:3 * nr_mix])
    x = torch.reshape(x, xs + [1]) + torch.zeros(xs + [nr_mix]).to(x.device)  # here and below: getting the means and adjusting them based on preceding sub-pixels
    m2 = torch.reshape(means[:, :, :, 1, :] + coeffs[:, :, :, 0, :] * x[:, :, :, 0, :], [xs[0], xs[1], xs[2], 1, nr_mix])
    m3 = torch.reshape(means[:, :, :, 2, :] + coeffs[:, :, :, 1, :] * x[:, :, :, 0, :] + coeffs[:, :, :, 2, :] * x[:, :, :, 1, :], [xs[0], xs[1], xs[2], 1, nr_mix])
    means = torch.cat([torch.reshape(means[:, :, :, 0, :], [xs[0], xs[1], xs[2], 1, nr_mix]), m2, m3], dim=3)
    centered_x = x - means
    inv_stdv = torch.exp(-log_scales)
    if low_bit:
        plus_in = inv_stdv * (centered_x + 1. / 31.)
        cdf_plus = torch.sigmoid(plus_in)
        min_in = inv_stdv * (centered_x - 1. / 31.)
    else:
        plus_in = inv_stdv * (centered_x + 1. / 255.)
        cdf_plus = torch.sigmoid(plus_in)
        min_in = inv_stdv * (centered_x - 1. / 255.)
    cdf_min = torch.sigmoid(min_in)
    log_cdf_plus = plus_in - F.softplus(plus_in)  # log probability for edge case of 0 (before scaling)
    log_one_minus_cdf_min = -F.softplus(min_in)  # log probability for edge case of 255 (before scaling)
    cdf_delta = cdf_plus - cdf_min  # probability for all other cases
    mid_in = inv_stdv * centered_x
    log_pdf_mid = mid_in - log_scales - 2. * F.softplus(mid_in)  # log probability in the center of the bin, to be used in extreme cases (not actually used in our code)

    # now select the right output: left edge case, right edge case, normal case, extremely low prob case (doesn't actually happen for us)

    # this is what we are really doing, but using the robust version below for extreme cases in other applications and to avoid NaN issue with tf.select()
    # log_probs = tf.select(x < -0.999, log_cdf_plus, tf.select(x > 0.999, log_one_minus_cdf_min, tf.log(cdf_delta)))

    # robust version, that still works if probabilities are below 1e-5 (which never happens in our code)
    # tensorflow backpropagates through tf.select() by multiplying with zero instead of selecting: this requires use to use some ugly tricks to avoid potential NaNs
    # the 1e-12 in tf.maximum(cdf_delta, 1e-12) is never actually used as output, it's purely there to get around the tf.select() gradient issue
    # if the probability on a sub-pixel is below 1e-5, we use an approximation based on the assumption that the log-density is constant in the bin of the observed sub-pixel value
    if low_bit:
        log_probs = torch.where(x < -0.999,
                                log_cdf_plus,
                                torch.where(x > 0.999,
                                            log_one_minus_cdf_min,
                                            torch.where(cdf_delta > 1e-5,
                                                        torch.log(const_max(cdf_delta, 1e-12)),
                                                        log_pdf_mid - np.log(15.5))))
    else:
        log_probs = torch.where(x < -0.999,
                                log_cdf_plus,
                                torch.where(x > 0.999,
                                            log_one_minus_cdf_min,
                                            torch.where(cdf_delta > 1e-5,
                                                        torch.log(const_max(cdf_delta, 1e-12)),
                                                        log_pdf_mid - np.log(127.5))))
    log_probs = log_probs.sum(dim=3) + log_prob_from_logits(logit_probs)
    mixture_probs = torch.logsumexp(log_probs, -1)
    res = -1. * mixture_probs.sum(dim=[1, 2]) / np.prod(xs[1:])
    return res


def sample_from_discretized_mix_logistic(l, nr_mix, eps=None, u=None):
    ls = [s for s in l.shape]
    xs = ls[:-1] + [3]
    # unpack parameters
    logit_probs = l[:, :, :, :nr_mix]
    l = torch.reshape(l[:, :, :, nr_mix:], xs + [nr_mix * 3])
    # sample mixture indicator from softmax
    if eps is None:
        eps = torch.empty(logit_probs.shape, device=l.device).uniform_(1e-5, 1. - 1e-5)
    amax = torch.argmax(logit_probs - torch.log(-torch.log(eps)), dim=3)
    sel = F.one_hot(amax, num_classes=nr_mix).float()
    sel = torch.reshape(sel, xs[:-1] + [1, nr_mix])
    # select logistic parameters
    means = (l[:, :, :, :, :nr_mix] * sel).sum(dim=4)
    log_scales = const_max((l[:, :, :, :, nr_mix:nr_mix * 2] * sel).sum(dim=4), -7.)
    coeffs = (torch.tanh(l[:, :, :, :, nr_mix * 2:nr_mix * 3]) * sel).sum(dim=4)
    # sample from logistic & clip to interval
    # we don't actually round to the nearest 8bit value when sampling
    if u is None:
        u = torch.empty(means.shape, device=means.device).uniform_(1e-5, 1. - 1e-5)
    x = means + torch.exp(log_scales) * (torch.log(u) - torch.log(1. - u))
    x0 = const_min(const_max(x[:, :, :, 0], -1.), 1.)
    x1 = const_min(const_max(x[:, :, :, 1] + coeffs[:, :, :, 0] * x0, -1.), 1.)
    x2 = const_min(const_max(x[:, :, :, 2] + coeffs[:, :, :, 1] * x0 + coeffs[:, :, :, 2] * x1, -1.), 1.)
    return torch.cat([torch.reshape(x0, xs[:-1] + [1]), torch.reshape(x1, xs[:-1] + [1]), torch.reshape(x2, xs[:-1] + [1])], dim=3), eps, u


def backtrack(H, sampler, imle, preprocess_fn, data, logprint, training_step_imle):
    latents = torch.randn([data.shape[0], H.latent_dim], requires_grad=True, dtype=torch.float32, device='cuda')
    snoise = [torch.randn([data.shape[0], s.shape[1], s.shape[2], s.shape[3]], dtype=torch.float32, device='cuda') for s in sampler.snoise_tmp]

    if H.restore_latent_path:
        logprint('restoring latent path')
        latents = torch.tensor(torch.load(f'{H.restore_latent_path}/latent-best.npy'), requires_grad=True, dtype=torch.float32, device='cuda')
        snoise = [torch.tensor(torch.load(f'{H.restore_latent_path}/snoise-best-{s.shape[2]}.npy'), requires_grad=True, dtype=torch.float32, device='cuda') for s in sampler.snoise_tmp]

    latent_optimizer = AdamW([latents], lr=H.latent_lr)
    if H.space == 'w':
        latent_optimizer = AdamW([latents] + snoise, lr=H.latent_lr)
    # latent_optimizer = SGD([latents] + snoise, lr=H.latent_lr)
    dists = torch.empty([data.shape[0]], dtype=torch.float32).cuda()

    sampler.calc_dists_existing(data, imle, dists=dists, latents=latents, snoise=snoise)
    print(f'initial dists: {dists.mean()}')

    best_loss = np.inf
    num_iters = 0

    while num_iters < H.reconstruct_iter_num:
        comb_dataset = ZippedDataset(data, TensorDataset(latents))
        data_loader = DataLoader(comb_dataset, batch_size=H.n_batch)
        for cur, indices in data_loader:
            x = cur
            lat = cur[1][0]
            _, target = preprocess_fn(x)
            cur_snoise = [s[indices] for s in snoise]
            training_step_imle(H, target.shape[0], target, lat, cur_snoise, imle, None, latent_optimizer, sampler.calc_loss)
            latents.grad.zero_()
            [s.grad.zero_() for s in snoise]
        num_iters += len(data)

        logprint(f'iteration: {num_iters}')
        # torch.save(latents.detach(), f'{H.save_dir}/latent-latest.npy')
        # for s in snoise:
        #     torch.save(s.detach(), f'{H.save_dir}/snoise-latest-{s.shape[2]}.npy')

        sampler.calc_dists_existing(data, imle, dists=dists, latents=latents, snoise=snoise)
        cur_mean = dists.mean()
        logprint(f'cur mean: {cur_mean}, best: {best_loss}')
        if cur_mean < best_loss:
            torch.save(latents.detach(), f'{H.save_dir}/latent-best.npy')
            for s in snoise:
                torch.save(s.detach(), f'{H.save_dir}/snoise-best-{s.shape[2]}.npy')
            logprint(f'improved: {cur_mean}')
            best_loss = cur_mean
            for i in range(data.shape[0]):
                samp = sampler.sample(latents[i:i+1], imle, [s[i:i+1] for s in snoise])
                imageio.imwrite(f'{H.save_dir}/{i}.png', samp[0])
                imageio.imwrite(f'{H.save_dir}/{i}-real.png', data[i])

        if num_iters >= H.reconstruct_iter_num:
            break


def reconstruct(H, sampler, imle, preprocess_fn, images, latents, snoise, name, logprint, training_step_imle):
    latent_optimizer = AdamW([latents], lr=H.latent_lr)
    generate_for_NN(sampler, images, latents.detach(), snoise, images.shape, imle,
                    f'{H.save_dir}/{name}-initial.png', logprint)
    for i in range(H.latent_epoch):
        for iter in range(H.reconstruct_iter_num):
            _, target = preprocess_fn([images])
            stat = training_step_imle(H, target.shape[0], target, latents, snoise, imle, None, latent_optimizer, sampler.calc_loss)

            latents.grad.zero_()
            if iter % 50 == 0:
                print('loss is: ', stat['loss'])
                generate_for_NN(sampler, images, latents.detach(), snoise, images.shape, imle,
                                f'{H.save_dir}/{name}-{iter}.png', logprint)

                torch.save(latents.detach(), '{}/reconstruct-latest.npy'.format(H.save_dir))