<commit_before>'''Simple GAN model

'''

import logging

import torch
from torch import autograd
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

#from resnets import ResEncoder as Discriminator
#from resnets import ResDecoder as Generator

from conv_decoders import SimpleConvDecoder as Generator
from convnets import SimpleConvEncoder as Discriminator

logger = logging.getLogger('cortex.models' + __name__)

GLOBALS = {'DIM_X': None, 'DIM_Y': None, 'DIM_C': None, 'DIM_L': None, 'DIM_Z': None}

#discriminator_args_ = dict(dim_h=64, batch_norm=True, f_size=3, n_steps=3)
#generator_args_ = dict(dim_h=64, batch_norm=True, f_size=3, n_steps=3)

discriminator_args_ = dict(dim_h=64, batch_norm=True, f_size=5, pad=2, stride=2, min_dim=7,
                           nonlinearity='LeakyReLU')
generator_args_ = dict(dim_h=64, batch_norm=True, f_size=5, pad=2, stride=1, n_steps=2)

#discriminator_args_ = dict(dim_h=64, batch_norm=True, min_dim=4, nonlinearity='LeakyReLU')
#generator_args_ = dict(dim_h=64, batch_norm=True, n_steps=3)

DEFAULTS = dict(
    data=dict(batch_size=64,
              noise_variables=dict(z=('normal', 64)),
              test_batch_size=64),
    optimizer=dict(
        optimizer='Adam',
        learning_rate=1e-4,
    ),
    model=dict(discriminator_args=discriminator_args_, generator_args=generator_args_),
    procedures=dict(measure='proxy_gan', penalty=False, boundary_seek=False),
    train=dict(
        epochs=200,
        summary_updates=100,
        archive_every=10
    )
)


def f_divergence(measure, real_out, fake_out, boundary_seek=False):
    if measure in ('gan', 'proxy_gan'):
        r = -F.softplus(-real_out)
        f = F.softplus(-fake_out) + fake_out
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'jsd':
        r = torch.log(2.) - F.softplus(-real_out)
        f = F.softplus(-fake_out) + fake_out + torch.log(2.)
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'xs':
        r = real_out ** 2
        f = -0.5 * ((torch.sqrt(fake_out ** 2) + 1.) ** 2)
        w = 0.5 * (1. - 1. / torch.sqrt(fake_out ** 2))
        b = (fake_out / 2.) ** 2

    elif measure == 'kl':
        r = real_out + 1.
        f = torch.exp(fake_out)
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'rkl':
        r = -torch.exp(-real_out)
        f = fake_out - 1.
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'dv':
        r = real_out
        f = torch.log(torch.exp(fake_out))
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'sh':
        r = 1. - torch.exp(-real_out)
        f = torch.exp(fake_out) - 1.
        w = torch.exp(fake_out)
        b = fake_out ** 2

    else:
        raise NotImplementedError(measure)
    d_loss = f.mean() - r.mean()

    if boundary_seek:
        g_loss = torch.mean(b)
    elif measure == 'proxy_gan':
        g_loss = torch.mean(F.softplus(-fake_out))
    else:
        g_loss = -torch.mean(f)

    return d_loss, g_loss, r, f, w, b


def apply_penalty(inputs, discriminator, generator, measure):
    Z = inputs['z']
    X = inputs['images']

    gen_out = generator(Z, nonlinearity=F.tanh)

    real = Variable(X.data.cuda(), requires_grad=True)
    fake = Variable(gen_out.data.cuda(), requires_grad=True)
    real_out = discriminator(real)
    fake_out = discriminator(fake)

    g_r = autograd.grad(outputs=real_out, inputs=real, grad_outputs=torch.ones(real_out.size()).cuda(),
                        create_graph=True, retain_graph=True, only_inputs=True)[0]

    g_f = autograd.grad(outputs=fake_out, inputs=fake, grad_outputs=torch.ones(fake_out.size()).cuda(),
                        create_graph=True, retain_graph=True, only_inputs=True)[0]

    if measure in ('gan', 'proxy_gan', 'jsd'):
        g_p = (0.5 * ((1. - F.sigmoid(real_out)) ** 2 * (g_r ** 2).sum(1).sum(1).sum(1)
                      + F.sigmoid(fake_out) ** 2 * (g_f ** 2).sum(1).sum(1).sum(1)))

    else:
        g_p = (0.5 * ((g_r ** 2).sum(1).sum(1).sum(1) + (g_f ** 2).sum(1).sum(1).sum(1)))

    g_p = torch.mean(g_p)
    return g_p


def gan(nets, inputs, measure=None, boundary_seek=False, penalty=None):
    Z = inputs['z']
    X = inputs['images']

    discriminator = nets['discriminator']
    generator = nets['generator']
    gen_out = generator(Z, nonlinearity=F.tanh)

    real_out = discriminator(X)
    fake_out = discriminator(gen_out)

    d_loss, g_loss, r, f, w, b = f_divergence(measure, real_out, fake_out, boundary_seek=boundary_seek)

    results = dict(g_loss=g_loss.data[0], d_loss=d_loss.data[0], boundary=torch.mean(b).data[0],
                   real=torch.mean(r).data[0], fake=torch.mean(f).data[0], w=torch.mean(w).data[0])
    samples = dict(images=dict(generated=0.5 * (gen_out.data + 1.), real=0.5 * (inputs['images'].data + 1.)))
    
    if penalty:
        p_term = apply_penalty(inputs, discriminator, generator, measure)

        d_loss += penalty * torch.mean(p_term)
        results['gradient penalty'] = torch.mean(p_term).data[0]

    return dict(generator=g_loss, discriminator=d_loss), results, samples, 'boundary'


def build_model(loss=None, discriminator_args=None, generator_args=None):
    discriminator_args = discriminator_args or {}
    generator_args = generator_args or {}

    shape = (DIM_X, DIM_Y, DIM_C)

    discriminator = Discriminator(shape, dim_out=1, **discriminator_args)
    generator = Generator(shape, dim_in=64, **generator_args)
    logger.debug(discriminator)
    logger.debug(generator)

    return dict(generator=generator, discriminator=discriminator), gan

<commit_msg>fixed jsd<commit_after>'''Simple GAN model

'''

import logging
import math

import torch
from torch import autograd
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

#from resnets import ResEncoder as Discriminator
#from resnets import ResDecoder as Generator

from conv_decoders import SimpleConvDecoder as Generator
from convnets import SimpleConvEncoder as Discriminator

logger = logging.getLogger('cortex.models' + __name__)

GLOBALS = {'DIM_X': None, 'DIM_Y': None, 'DIM_C': None, 'DIM_L': None, 'DIM_Z': None}

#discriminator_args_ = dict(dim_h=64, batch_norm=True, f_size=3, n_steps=3)
#generator_args_ = dict(dim_h=64, batch_norm=True, f_size=3, n_steps=3)

discriminator_args_ = dict(dim_h=64, batch_norm=True, f_size=5, pad=2, stride=2, min_dim=7,
                           nonlinearity='LeakyReLU')
generator_args_ = dict(dim_h=64, batch_norm=True, f_size=5, pad=2, stride=1, n_steps=2)

#discriminator_args_ = dict(dim_h=64, batch_norm=True, min_dim=4, nonlinearity='LeakyReLU')
#generator_args_ = dict(dim_h=64, batch_norm=True, n_steps=3)

DEFAULTS = dict(
    data=dict(batch_size=64,
              noise_variables=dict(z=('normal', 64)),
              test_batch_size=64),
    optimizer=dict(
        optimizer='Adam',
        learning_rate=1e-4,
    ),
    model=dict(discriminator_args=discriminator_args_, generator_args=generator_args_),
    procedures=dict(measure='proxy_gan', penalty=False, boundary_seek=False),
    train=dict(
        epochs=200,
        summary_updates=100,
        archive_every=10
    )
)


def f_divergence(measure, real_out, fake_out, boundary_seek=False):
    log_2 = math.log(2.)

    if measure in ('gan', 'proxy_gan'):
        r = -F.softplus(-real_out)
        f = F.softplus(-fake_out) + fake_out
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'jsd':
        r = log_2 - F.softplus(-real_out)
        f = F.softplus(-fake_out) + fake_out + log_2
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'xs':
        r = real_out ** 2
        f = -0.5 * ((torch.sqrt(fake_out ** 2) + 1.) ** 2)
        w = 0.5 * (1. - 1. / torch.sqrt(fake_out ** 2))
        b = (fake_out / 2.) ** 2

    elif measure == 'kl':
        r = real_out + 1.
        f = torch.exp(fake_out)
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'rkl':
        r = -torch.exp(-real_out)
        f = fake_out - 1.
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'dv':
        r = real_out
        f = torch.log(torch.exp(fake_out))
        w = torch.exp(fake_out)
        b = fake_out ** 2

    elif measure == 'sh':
        r = 1. - torch.exp(-real_out)
        f = torch.exp(fake_out) - 1.
        w = torch.exp(fake_out)
        b = fake_out ** 2

    else:
        raise NotImplementedError(measure)
    d_loss = f.mean() - r.mean()

    if boundary_seek:
        g_loss = torch.mean(b)
    elif measure == 'proxy_gan':
        g_loss = torch.mean(F.softplus(-fake_out))
    else:
        g_loss = -torch.mean(f)

    return d_loss, g_loss, r, f, w, b


def apply_penalty(inputs, discriminator, generator, measure):
    Z = inputs['z']
    X = inputs['images']

    gen_out = generator(Z, nonlinearity=F.tanh)

    real = Variable(X.data.cuda(), requires_grad=True)
    fake = Variable(gen_out.data.cuda(), requires_grad=True)
    real_out = discriminator(real)
    fake_out = discriminator(fake)

    g_r = autograd.grad(outputs=real_out, inputs=real, grad_outputs=torch.ones(real_out.size()).cuda(),
                        create_graph=True, retain_graph=True, only_inputs=True)[0]

    g_f = autograd.grad(outputs=fake_out, inputs=fake, grad_outputs=torch.ones(fake_out.size()).cuda(),
                        create_graph=True, retain_graph=True, only_inputs=True)[0]

    if measure in ('gan', 'proxy_gan', 'jsd'):
        g_p = (0.5 * ((1. - F.sigmoid(real_out)) ** 2 * (g_r ** 2).sum(1).sum(1).sum(1)
                      + F.sigmoid(fake_out) ** 2 * (g_f ** 2).sum(1).sum(1).sum(1)))

    else:
        g_p = (0.5 * ((g_r ** 2).sum(1).sum(1).sum(1) + (g_f ** 2).sum(1).sum(1).sum(1)))

    g_p = torch.mean(g_p)
    return g_p


def gan(nets, inputs, measure=None, boundary_seek=False, penalty=None):
    Z = inputs['z']
    X = inputs['images']

    discriminator = nets['discriminator']
    generator = nets['generator']
    gen_out = generator(Z, nonlinearity=F.tanh)

    real_out = discriminator(X)
    fake_out = discriminator(gen_out)

    d_loss, g_loss, r, f, w, b = f_divergence(measure, real_out, fake_out, boundary_seek=boundary_seek)

    results = dict(g_loss=g_loss.data[0], d_loss=d_loss.data[0], boundary=torch.mean(b).data[0],
                   real=torch.mean(r).data[0], fake=torch.mean(f).data[0], w=torch.mean(w).data[0])
    samples = dict(images=dict(generated=0.5 * (gen_out.data + 1.), real=0.5 * (inputs['images'].data + 1.)))

    if penalty:
        p_term = apply_penalty(inputs, discriminator, generator, measure)

        d_loss += penalty * torch.mean(p_term)
        results['gradient penalty'] = torch.mean(p_term).data[0]

    return dict(generator=g_loss, discriminator=d_loss), results, samples, 'boundary'


def build_model(loss=None, discriminator_args=None, generator_args=None):
    discriminator_args = discriminator_args or {}
    generator_args = generator_args or {}

    shape = (DIM_X, DIM_Y, DIM_C)

    discriminator = Discriminator(shape, dim_out=1, **discriminator_args)
    generator = Generator(shape, dim_in=64, **generator_args)
    logger.debug(discriminator)
    logger.debug(generator)

    return dict(generator=generator, discriminator=discriminator), gan


{'fraction_non_alphanumeric': 0.14339523475823407, 'fraction_numerical': 0.019096005606166783, 'mean_word_length': 3.0921146953405017, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2586072', 'n_tokens_mistral': 4276, 'n_tokens_neox': 4110, 'n_words': 811}