<gh_stars>0
from __future__ import print_function
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
import torch.nn.modules.loss
from torch.nn.modules.loss import _WeightedLoss
from torch.nn.modules.loss import BCELoss
from torch.nn import functional as F
import torchvision.transforms.functional as TVF
import torchvision.models as models
from torch.autograd import Variable
import faiss
import numpy as np
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data import TensorDataset, DataLoader

DIMENSION_REDUCTION_NN_INPUT_DIM = 224
DIMENSION_REDUCTION_NN_OUTPUT_DIM = 512
COMPRESSED_KNN_INDEX_SIZE = 100
NN_K = 10
N_TRAINING_SAMPLES = 64000
#N_TRAINING_SAMPLES = 640

#### Define what we need to map the training images into high dimensional space

#Load the pretrained model
pretrained_image_classification_model = models.resnet18(pretrained=True)

# Use the model object to select the desired layer
image_dim_reduction_nn = nn.Sequential(*list(pretrained_image_classification_model.children())[:-1])

# Set model to evaluation mode
image_dim_reduction_nn.eval()
image_dim_reduction_nn.cuda()


parser = argparse.ArgumentParser()
parser.add_argument('--dataroot', required=True, help='path to dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
parser.add_argument('--batchSize', type=int, default=64, help='input batch size')
parser.add_argument('--imageSize', type=int, default=64, help='the height / width of the input image to network')
parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--niter', type=int, default=25, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--ngpu', type=int, default=1, help='number of GPUs to use')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
parser.add_argument('--manualSeed', type=int, help='manual seed')
parser.add_argument('--negan', action='store_true', help='biases search toward underexploited search space')

opt = parser.parse_args()
print(opt)

try:
    os.makedirs(opt.outf)
except OSError:
    pass

if opt.manualSeed is None:
    opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
    torch.cuda.manual_seed_all(opt.manualSeed)

cudnn.benchmark = True

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

train_sampler = SubsetRandomSampler(np.arange(N_TRAINING_SAMPLES, dtype=np.int64))

dataset = dset.LSUN(opt.dataroot, classes=['bedroom_train'],
                    transform=transforms.Compose([
                        transforms.Resize(opt.imageSize),
                        transforms.CenterCrop(opt.imageSize),
                        transforms.ToTensor(),
                        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                    ]))

dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize,
                                         shuffle=False, num_workers=int(opt.workers), sampler=train_sampler)


# Load the training set again, scaled and noramlized for the resnet model
dataset_dim_reduction_baseline = dset.LSUN(opt.dataroot, classes=['bedroom_train'],
                    transform=transforms.Compose([
                        transforms.Resize(DIMENSION_REDUCTION_NN_INPUT_DIM),
                        transforms.CenterCrop(DIMENSION_REDUCTION_NN_INPUT_DIM),
                        transforms.ToTensor(),
                        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
                    ]))

dataloader_dim_reduction_baseline = torch.utils.data.DataLoader(dataset_dim_reduction_baseline, batch_size=opt.batchSize,
                                         shuffle=False, num_workers=int(opt.workers), sampler=train_sampler)


res = faiss.StandardGpuResources()
index_flat = faiss.IndexFlatL2(DIMENSION_REDUCTION_NN_OUTPUT_DIM)
# index_ivf = faiss.IndexIVFFlat(quantizer, DIMENSION_REDUCTION_NN_OUTPUT_DIM, COMPRESSED_KNN_INDEX_SIZE, faiss.METRIC_L2)
# gpu_index = faiss.index_cpu_to_gpu(res, 0, index_ivf)
gpu_index = faiss.index_cpu_to_gpu(res, 0, index_flat)

baseline_np_array = np.ones((N_TRAINING_SAMPLES,DIMENSION_REDUCTION_NN_OUTPUT_DIM),np.float32)

baseline_image_input = torch.FloatTensor(opt.batchSize, 3, DIMENSION_REDUCTION_NN_INPUT_DIM, DIMENSION_REDUCTION_NN_INPUT_DIM)
if opt.cuda:
    baseline_image_input = baseline_image_input.cuda()

## Iterate through the training data, computing vectors to use in sparseness calculations
for i, baseline_image_batch in enumerate(dataloader_dim_reduction_baseline, 0):
    real_cpu, _ = baseline_image_batch
    batch_size = real_cpu.size(0)
    if opt.cuda:
        real_cpu = real_cpu.cuda()

    baseline_image_input.resize_as_(real_cpu).copy_(real_cpu)

    baseline_np_array[range(i*opt.batchSize,(i+1)*opt.batchSize),:] = image_dim_reduction_nn(baseline_image_input)[:,:,0,0].detach().cpu().numpy()

    print("Generating baseline batch %d" % (i))


print(np.amin(baseline_np_array), np.amax(baseline_np_array), np.mean(baseline_np_array), np.std(baseline_np_array))

# assert not gpu_index_ivf.is_trained
# gpu_index_ivf.train(baseline_np_array)
# assert gpu_index_ivf.is_trained

gpu_index.add(baseline_np_array)
print(gpu_index.ntotal)


def weights_from_reduced_dim_images(images):
    # Scale images to be compatible with dim reduction NN
    scaler = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize([224,224]),
        transforms.ToTensor()
    ])

    scaled_images = torch.FloatTensor(len(images), 3, DIMENSION_REDUCTION_NN_INPUT_DIM, DIMENSION_REDUCTION_NN_INPUT_DIM)

    for i in range(0, len(images)):
        scaled_images[i,:,:,:] = scaler(images[i].detach().cpu())
        # TODO, do I have to convert to cpu ^
    
    # Get reduced dim vector rep of images
    reduced_dim_images = image_dim_reduction_nn(scaled_images.cuda())

    # Compute the sparseness metric for each reduced dim vector
    D, I = gpu_index.search(reduced_dim_images[:,:,0,0].detach().cpu().numpy(), NN_K)

    # print("Closest images = ")
    print(I)
    #print(D)
    reduced_dim_images_debug = reduced_dim_images[:,:,0,0].detach().cpu().numpy()
    print(np.amin(reduced_dim_images_debug), np.amax(reduced_dim_images_debug), np.mean(reduced_dim_images_debug), np.std(reduced_dim_images_debug))

    def mean_of_finite_distances(x):
        return np.mean(x[np.where( x < 1e6)])

    mean_distance = np.apply_along_axis(mean_of_finite_distances, axis=1, arr=D)

    return torch.from_numpy(mean_distance)


ngpu = int(opt.ngpu)
nz = int(opt.nz)
ngf = int(opt.ngf)
ndf = int(opt.ndf)
nc = 3


# custom weights initialization called on netG and netD
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)


class _netG(nn.Module):
    def __init__(self, ngpu):
        super(_netG, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),
            #nn.ConvTranspose2d(    ngf,      nc, 164, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
        else:
            output = self.main(input)
        return output


netG = _netG(ngpu)
netG.apply(weights_init)
if opt.netG != '':
    netG.load_state_dict(torch.load(opt.netG))
print(netG)


class _netD(nn.Module):
    def __init__(self, ngpu):
        super(_netD, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*8) x 4 x 4
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            #nn.Conv2d(ndf * 8, 1, 14, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
        else:
            output = self.main(input)
        return output.view(-1, 1).squeeze(1)


netD = _netD(ngpu)
netD.apply(weights_init)
if opt.netD != '':
    netD.load_state_dict(torch.load(opt.netD))
print(netD)

criterionD = BCELoss()
criterionG = BCELoss()

input = torch.FloatTensor(opt.batchSize, 3, opt.imageSize, opt.imageSize)
noise = torch.FloatTensor(opt.batchSize, nz, 1, 1)
fixed_noise = torch.FloatTensor(opt.batchSize, nz, 1, 1).normal_(0, 1)
label = torch.FloatTensor(opt.batchSize)
real_label = 1
fake_label = 0

if opt.cuda:
    netD.cuda()
    netG.cuda()
    criterionG.cuda()
    criterionD.cuda()
    input, label = input.cuda(), label.cuda()
    noise, fixed_noise = noise.cuda(), fixed_noise.cuda()

fixed_noise = Variable(fixed_noise)

# setup optimizer
optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

for epoch in range(opt.niter):
    for i, data in enumerate(dataloader, 0):
        ############################
        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
        ###########################
        # train with real
        netD.zero_grad()
        real_cpu, _ = data
        batch_size = real_cpu.size(0)
        if opt.cuda:
            real_cpu = real_cpu.cuda()
        
        input.resize_as_(real_cpu).copy_(real_cpu)
        label.resize_(batch_size).fill_(real_label)
        inputv = Variable(input)
        labelv = Variable(label)

        output = netD(inputv)

        errD_real = criterionD(output, labelv)
        errD_real.backward()
        D_x = output.data.mean()

        # train with fake
        noise.resize_(batch_size, nz, 1, 1).normal_(0, 1)
        noisev = Variable(noise)
        fake = netG(noisev)
        labelv = Variable(label.fill_(fake_label))
        output = netD(fake.detach())
        errD_fake = criterionD(output, labelv)
        errD_fake.backward()
        D_G_z1 = output.data.mean()
        errD = errD_real + errD_fake
        optimizerD.step()

        ############################
        # (2) Update G network: maximize log(D(G(z)))
        ###########################
        netG.zero_grad()
        labelv = Variable(label.fill_(real_label))  # fake labels are real for generator cost
        output = netD(fake)
        errG = criterionG(output, labelv)

        # Set the weights on the criteria to correspond to density of the image space
        weights = weights_from_reduced_dim_images(fake)
        if opt.negan and epoch > 25:
            print("Setting weights...")
            errG.weights = weights

        errG.backward()
        D_G_z2 = output.data.mean()
        optimizerG.step()

        print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f Avg_Dist_From_Training: %.4f'
              % (epoch, opt.niter, i, len(dataloader),
                 errD.data[0], errG.data[0], D_x, D_G_z1, D_G_z2, weights.mean()))
        if i % 100 == 0:
            vutils.save_image(real_cpu,
                    '%s/real_samples.png' % opt.outf,
                    normalize=True)
            fake = netG(fixed_noise)
            vutils.save_image(fake.data,
                    '%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch),
                    normalize=True)

    # do checkpointing
    torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch))
    torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (opt.outf, epoch))

from __future__ import print_function
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
import torch.nn.modules.loss
from torch.nn.modules.loss import _WeightedLoss
from torch.nn.modules.loss import BCELoss
from torch.nn import functional as F
import torchvision.transforms.functional as TVF
import torchvision.models as models
from torch.autograd import Variable
import faiss
import numpy as np
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data import TensorDataset, DataLoader

DIMENSION_REDUCTION_NN_INPUT_DIM = 224
DIMENSION_REDUCTION_NN_OUTPUT_DIM = 512
COMPRESSED_KNN_INDEX_SIZE = 100
NN_K = 10
N_TRAINING_SAMPLES = 64000
#N_TRAINING_SAMPLES = 640

#### Define what we need to map the training images into high dimensional space

#Load the pretrained model
pretrained_image_classification_model = models.resnet18(pretrained=True)

# Use the model object to select the desired layer
image_dim_reduction_nn = nn.Sequential(*list(pretrained_image_classification_model.children())[:-1])

# Set model to evaluation mode
image_dim_reduction_nn.eval()
image_dim_reduction_nn.cuda()


parser = argparse.ArgumentParser()
parser.add_argument('--dataroot', required=True, help='path to dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
parser.add_argument('--batchSize', type=int, default=64, help='input batch size')
parser.add_argument('--imageSize', type=int, default=64, help='the height / width of the input image to network')
parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--niter', type=int, default=25, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--ngpu', type=int, default=1, help='number of GPUs to use')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
parser.add_argument('--manualSeed', type=int, help='manual seed')
parser.add_argument('--negan', action='store_true', help='biases search toward underexploited search space')

opt = parser.parse_args()
print(opt)

try:
    os.makedirs(opt.outf)
except OSError:
    pass

if opt.manualSeed is None:
    opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
    torch.cuda.manual_seed_all(opt.manualSeed)

cudnn.benchmark = True

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

train_sampler = SubsetRandomSampler(np.arange(N_TRAINING_SAMPLES, dtype=np.int64))

dataset = dset.LSUN(opt.dataroot, classes=['bedroom_train'],
                    transform=transforms.Compose([
                        transforms.Resize(opt.imageSize),
                        transforms.CenterCrop(opt.imageSize),
                        transforms.ToTensor(),
                        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                    ]))

dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize,
                                         shuffle=False, num_workers=int(opt.workers), sampler=train_sampler)


# Load the training set again, scaled and noramlized for the resnet model
dataset_dim_reduction_baseline = dset.LSUN(opt.dataroot, classes=['bedroom_train'],
                    transform=transforms.Compose([
                        transforms.Resize(DIMENSION_REDUCTION_NN_INPUT_DIM),
                        transforms.CenterCrop(DIMENSION_REDUCTION_NN_INPUT_DIM),
                        transforms.ToTensor(),
                        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
                    ]))

dataloader_dim_reduction_baseline = torch.utils.data.DataLoader(dataset_dim_reduction_baseline, batch_size=opt.batchSize,
                                         shuffle=False, num_workers=int(opt.workers), sampler=train_sampler)


res = faiss.StandardGpuResources()
index_flat = faiss.IndexFlatL2(DIMENSION_REDUCTION_NN_OUTPUT_DIM)
# index_ivf = faiss.IndexIVFFlat(quantizer, DIMENSION_REDUCTION_NN_OUTPUT_DIM, COMPRESSED_KNN_INDEX_SIZE, faiss.METRIC_L2)
# gpu_index = faiss.index_cpu_to_gpu(res, 0, index_ivf)
gpu_index = faiss.index_cpu_to_gpu(res, 0, index_flat)

baseline_np_array = np.ones((N_TRAINING_SAMPLES,DIMENSION_REDUCTION_NN_OUTPUT_DIM),np.float32)

baseline_image_input = torch.FloatTensor(opt.batchSize, 3, DIMENSION_REDUCTION_NN_INPUT_DIM, DIMENSION_REDUCTION_NN_INPUT_DIM)
if opt.cuda:
    baseline_image_input = baseline_image_input.cuda()

## Iterate through the training data, computing vectors to use in sparseness calculations
for i, baseline_image_batch in enumerate(dataloader_dim_reduction_baseline, 0):
    real_cpu, _ = baseline_image_batch
    batch_size = real_cpu.size(0)
    if opt.cuda:
        real_cpu = real_cpu.cuda()

    baseline_image_input.resize_as_(real_cpu).copy_(real_cpu)

    baseline_np_array[range(i*opt.batchSize,(i+1)*opt.batchSize),:] = image_dim_reduction_nn(baseline_image_input)[:,:,0,0].detach().cpu().numpy()

    print("Generating baseline batch %d" % (i))


print(np.amin(baseline_np_array), np.amax(baseline_np_array), np.mean(baseline_np_array), np.std(baseline_np_array))

# assert not gpu_index_ivf.is_trained
# gpu_index_ivf.train(baseline_np_array)
# assert gpu_index_ivf.is_trained

gpu_index.add(baseline_np_array)
print(gpu_index.ntotal)


def weights_from_reduced_dim_images(images):
    # Scale images to be compatible with dim reduction NN
    scaler = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize([224,224]),
        transforms.ToTensor()
    ])

    scaled_images = torch.FloatTensor(len(images), 3, DIMENSION_REDUCTION_NN_INPUT_DIM, DIMENSION_REDUCTION_NN_INPUT_DIM)

    for i in range(0, len(images)):
        scaled_images[i,:,:,:] = scaler(images[i].detach().cpu())
        # TODO, do I have to convert to cpu ^
    
    # Get reduced dim vector rep of images
    reduced_dim_images = image_dim_reduction_nn(scaled_images.cuda())

    # Compute the sparseness metric for each reduced dim vector
    D, I = gpu_index.search(reduced_dim_images[:,:,0,0].detach().cpu().numpy(), NN_K)

    # print("Closest images = ")
    print(I)
    #print(D)
    reduced_dim_images_debug = reduced_dim_images[:,:,0,0].detach().cpu().numpy()
    print(np.amin(reduced_dim_images_debug), np.amax(reduced_dim_images_debug), np.mean(reduced_dim_images_debug), np.std(reduced_dim_images_debug))

    def mean_of_finite_distances(x):
        return np.mean(x[np.where( x < 1e6)])

    mean_distance = np.apply_along_axis(mean_of_finite_distances, axis=1, arr=D)

    return torch.from_numpy(mean_distance)


ngpu = int(opt.ngpu)
nz = int(opt.nz)
ngf = int(opt.ngf)
ndf = int(opt.ndf)
nc = 3


# custom weights initialization called on netG and netD
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)


class _netG(nn.Module):
    def __init__(self, ngpu):
        super(_netG, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),
            #nn.ConvTranspose2d(    ngf,      nc, 164, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
        else:
            output = self.main(input)
        return output


netG = _netG(ngpu)
netG.apply(weights_init)
if opt.netG != '':
    netG.load_state_dict(torch.load(opt.netG))
print(netG)


class _netD(nn.Module):
    def __init__(self, ngpu):
        super(_netD, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*8) x 4 x 4
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            #nn.Conv2d(ndf * 8, 1, 14, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
        else:
            output = self.main(input)
        return output.view(-1, 1).squeeze(1)


netD = _netD(ngpu)
netD.apply(weights_init)
if opt.netD != '':
    netD.load_state_dict(torch.load(opt.netD))
print(netD)

criterionD = BCELoss()
criterionG = BCELoss()

input = torch.FloatTensor(opt.batchSize, 3, opt.imageSize, opt.imageSize)
noise = torch.FloatTensor(opt.batchSize, nz, 1, 1)
fixed_noise = torch.FloatTensor(opt.batchSize, nz, 1, 1).normal_(0, 1)
label = torch.FloatTensor(opt.batchSize)
real_label = 1
fake_label = 0

if opt.cuda:
    netD.cuda()
    netG.cuda()
    criterionG.cuda()
    criterionD.cuda()
    input, label = input.cuda(), label.cuda()
    noise, fixed_noise = noise.cuda(), fixed_noise.cuda()

fixed_noise = Variable(fixed_noise)

# setup optimizer
optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

for epoch in range(opt.niter):
    for i, data in enumerate(dataloader, 0):
        ############################
        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
        ###########################
        # train with real
        netD.zero_grad()
        real_cpu, _ = data
        batch_size = real_cpu.size(0)
        if opt.cuda:
            real_cpu = real_cpu.cuda()
        
        input.resize_as_(real_cpu).copy_(real_cpu)
        label.resize_(batch_size).fill_(real_label)
        inputv = Variable(input)
        labelv = Variable(label)

        output = netD(inputv)

        errD_real = criterionD(output, labelv)
        errD_real.backward()
        D_x = output.data.mean()

        # train with fake
        noise.resize_(batch_size, nz, 1, 1).normal_(0, 1)
        noisev = Variable(noise)
        fake = netG(noisev)
        labelv = Variable(label.fill_(fake_label))
        output = netD(fake.detach())
        errD_fake = criterionD(output, labelv)
        errD_fake.backward()
        D_G_z1 = output.data.mean()
        errD = errD_real + errD_fake
        optimizerD.step()

        ############################
        # (2) Update G network: maximize log(D(G(z)))
        ###########################
        netG.zero_grad()
        labelv = Variable(label.fill_(real_label))  # fake labels are real for generator cost
        output = netD(fake)
        errG = criterionG(output, labelv)

        # Set the weights on the criteria to correspond to density of the image space
        weights = weights_from_reduced_dim_images(fake)
        if opt.negan and epoch > 25:
            print("Setting weights...")
            errG.weights = weights

        errG.backward()
        D_G_z2 = output.data.mean()
        optimizerG.step()

        print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f Avg_Dist_From_Training: %.4f'
              % (epoch, opt.niter, i, len(dataloader),
                 errD.data[0], errG.data[0], D_x, D_G_z1, D_G_z2, weights.mean()))
        if i % 100 == 0:
            vutils.save_image(real_cpu,
                    '%s/real_samples.png' % opt.outf,
                    normalize=True)
            fake = netG(fixed_noise)
            vutils.save_image(fake.data,
                    '%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch),
                    normalize=True)

    # do checkpointing
    torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch))
    torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (opt.outf, epoch))
#N_TRAINING_SAMPLES = 640 #### Define what we need to map the training images into high dimensional space #Load the pretrained model # Use the model object to select the desired layer # Set model to evaluation mode # Load the training set again, scaled and noramlized for the resnet model # index_ivf = faiss.IndexIVFFlat(quantizer, DIMENSION_REDUCTION_NN_OUTPUT_DIM, COMPRESSED_KNN_INDEX_SIZE, faiss.METRIC_L2) # gpu_index = faiss.index_cpu_to_gpu(res, 0, index_ivf) ## Iterate through the training data, computing vectors to use in sparseness calculations # assert not gpu_index_ivf.is_trained # gpu_index_ivf.train(baseline_np_array) # assert gpu_index_ivf.is_trained # Scale images to be compatible with dim reduction NN # TODO, do I have to convert to cpu ^ # Get reduced dim vector rep of images # Compute the sparseness metric for each reduced dim vector # print("Closest images = ") #print(D) # custom weights initialization called on netG and netD # input is Z, going into a convolution # state size. (ngf*8) x 4 x 4 # state size. (ngf*4) x 8 x 8 # state size. (ngf*2) x 16 x 16 # state size. (ngf) x 32 x 32 #nn.ConvTranspose2d(    ngf,      nc, 164, 2, 1, bias=False), # state size. (nc) x 64 x 64 # input is (nc) x 64 x 64 # state size. (ndf) x 32 x 32 # state size. (ndf*2) x 16 x 16 # state size. (ndf*4) x 8 x 8 # state size. (ndf*8) x 4 x 4 #nn.Conv2d(ndf * 8, 1, 14, 1, 0, bias=False), # setup optimizer ############################ # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ########################### # train with real # train with fake ############################ # (2) Update G network: maximize log(D(G(z))) ########################### # fake labels are real for generator cost # Set the weights on the criteria to correspond to density of the image space # do checkpointing