text: import os, sys
  sys.path.append(os.getcwd())
  
  import torch
  import torchvision
  from torch import nn
  from torch import autograd
  from torch import optim
  import torch.nn.functional as F
  
  
  import time
  import tflib as lib
  import tflib.save_images
  import tflib.mnist
  import tflib.cifar10
  import tflib.plot
  #import tflib.inception_score
  
  import numpy as np
  
  # Download CIFAR-10 (Python version) at
  # https://www.cs.toronto.edu/~kriz/cifar.html and fill in the path to the
  # extracted files here!
  DATA_DIR = '/mnt/7FC1A7CD7234342C/cifar-10-batches-py/'
  if len(DATA_DIR) == 0:
      raise Exception('Please specify path to data directory in gan_cifar.py!')
  
  MODE = 'wgan-gp' # Valid options are dcgan, wgan, or wgan-gp
  DIM = 128 # This overfits substantially; you're probably better off with 64
  LAMBDA = 10 # Gradient penalty lambda hyperparameter
  CRITIC_ITERS = 5 # How many critic iterations per generator iteration
  BATCH_SIZE = 64 # Batch size
  ITERS = 100000 # How many generator iterations to train for
  OUTPUT_DIM = 3072 # Number of pixels in CIFAR10 (3*32*32)
  
  KERNEL_SIZE = 4
  
  def _upscale_resize(in_dim, out_dim, kernel_size):
      return nn.Sequential(
          nn.InstanceNorm2d(in_dim, affine=True),
          nn.Upsample(scale_factor=2, mode='nearest'),
          nn.ReflectionPad2d((1,2,1,2)),
          nn.Conv2d(in_dim, out_dim, kernel_size, bias=False)
      )
  
  class Generator(nn.Module):
      def __init__(self):
          super(Generator, self).__init__()
          preprocess = nn.Sequential(
              #nn.InstanceNorm2d(4 * 4 * 4 * DIM),
              nn.Linear(128, 4 * 4 * 4 * DIM),
              nn.ReLU(True),
          )
  
          block1 = nn.Sequential(
              nn.InstanceNorm2d(4 * DIM, affine=True),
              nn.ConvTranspose2d(4 * DIM, 2 * DIM, KERNEL_SIZE, stride=2, padding=1, bias=False),
              nn.ReLU(True),
          )
          block2 = nn.Sequential(
              nn.InstanceNorm2d(2 * DIM, affine=True),
              nn.ConvTranspose2d(2 * DIM, DIM, KERNEL_SIZE, stride=2, padding=1, bias=False),
              nn.ReLU(True),
          )
          #block1 = _upscale_resize(4 * DIM, 2 * DIM, KERNEL_SIZE)
          #block2 = _upscale_resize(2 * DIM, DIM, KERNEL_SIZE)
          self.last_norm  = nn.InstanceNorm2d(DIM, affine=True)
          deconv_out = nn.ConvTranspose2d(DIM, 3, KERNEL_SIZE, stride=2, padding=1, bias=False)
          self.out_norm = nn.InstanceNorm2d(3)
          
          self.preprocess = preprocess
          self.block1 = block1
          self.block2 = block2
          self.deconv_out = deconv_out
          self.tanh = nn.Tanh()
  
      def forward(self, input):
          output = self.preprocess(input)
          output = output.view(-1, 4 * DIM, 4, 4)
          #print(output.size())
          output = self.block1(output)
          #print(output.size())
          output = self.block2(output)
          #print(output.size())
          output = self.deconv_out(self.last_norm(output))
          #output = self.deconv_out(output)
          output = self.tanh(output)
          #output = self.out_norm(output)
          return output.view(-1, 3, 32, 32)
  
  
  class Discriminator(nn.Module):
      def __init__(self):
          super(Discriminator, self).__init__()
          main = nn.Sequential(
              nn.Conv2d(3, DIM, KERNEL_SIZE, 2, padding=1, bias=False),
              nn.InstanceNorm2d(DIM),
              nn.LeakyReLU(0.2, True),
              nn.Conv2d(DIM, 2 * DIM, KERNEL_SIZE, 2, padding=1, bias=False),
              nn.InstanceNorm2d(2 * DIM),
              nn.LeakyReLU(0.2, True),
              nn.Conv2d(2 * DIM, 4 * DIM, KERNEL_SIZE, 2, padding=1, bias=False),
              nn.InstanceNorm2d(4 * DIM),
              nn.LeakyReLU(0.2, True),
          )
  
          self.main = main
          self.linear = nn.Linear(4*4*4*DIM, 1)
  
      def forward(self, input):
          output = self.main(input)
          before_linear = output.view(-1, 4 * 4 * 4 * DIM)
          output = self.linear(before_linear)
          return output
  
  def critic_schedule():
      for i in range(10):
          yield 100
      while True:
          yield 5
  
  def gen_schedule():
      for i in range(10):
          yield 1
      for i in range(100):
          yield 1
      for i in range(7000):
          yield 1
      while True:
          yield 1
  
  # 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('Norm') != -1:
          if m.weight is not None:
              m.weight.data.normal_(1.0, 0.02)
              m.bias.data.fill_(0)
  
  netG = Generator()
  netD = Discriminator()
  netG.apply(weights_init)
  netD.apply(weights_init)
  print(netG)
  print(netD)
  use_cuda = torch.cuda.is_available()
  mse_loss = torch.nn.MSELoss()
  if use_cuda:
      gpu = 0
      # makes things slower?!
      torch.backends.cudnn.benchmark = True
  if use_cuda:
      netD = netD.cuda(gpu)
      netG = netG.cuda(gpu)
      mse_loss = mse_loss.cuda(gpu)
  one = torch.FloatTensor([1])
  mone = one * -1
  if use_cuda:
      one = one.cuda(gpu)
      mone = mone.cuda(gpu)
  
  optimizerD = optim.Adam(netD.parameters(), lr=1e-3, betas=(0.5, 0.9))
  optimizerG = optim.Adam(netG.parameters(), lr=1e-3, betas=(0.5, 0.9))
  
  netG.train()
  netD.train()
  
  def calc_gradient_penalty(netD, netG, real_data, fake_data, real_data_grad):
      # print "real_data: ", real_data.size(), fake_data.size()
      alpha = torch.rand(BATCH_SIZE, 1)
      alpha = alpha.expand(BATCH_SIZE, real_data.nelement()//BATCH_SIZE).contiguous().view(BATCH_SIZE, 3, 32, 32)
      alpha = alpha.cuda(gpu) if use_cuda else alpha
  
      interpolates = alpha * real_data + ((1 - alpha) * fake_data)
      #interpolates = real_data + (alpha * real_data_grad.data * torch.norm(real_data_grad.data, p=2, dim=1).unsqueeze(1))
  
      if use_cuda:
          interpolates = interpolates.cuda(gpu)
      interpolates = autograd.Variable(interpolates, requires_grad=True)
  
      disc_interpolates = netD(interpolates)
  
      gradients = autograd.grad(outputs=disc_interpolates, inputs=interpolates,
                                grad_outputs=torch.ones(disc_interpolates.size()).cuda(gpu) if use_cuda else torch.ones(
                                    disc_interpolates.size()),
                                create_graph=True, retain_graph=True, only_inputs=True)[0]
      gradients = gradients.view(gradients.size(0), -1)
  
      #gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * LAMBDA
      # one-sided penalty
      gradient_penalty = (F.relu(gradients.norm(2, dim=1) - 1, inplace=True) ** 2).mean() * LAMBDA
  
      return gradient_penalty
  
  # For generating samples
  def generate_image(frame, netG):
      fixed_noise_128 = torch.randn(128, 128)
      if use_cuda:
          fixed_noise_128 = fixed_noise_128.cuda(gpu)
      noisev = autograd.Variable(fixed_noise_128, volatile=True)
      netG.eval()
      samples = netG(noisev)
      netG.train()
      samples = samples.view(-1, 3, 32, 32)
      samples = samples.mul(0.5).add(0.5)
      samples = samples.cpu().data.numpy()
  
      lib.save_images.save_images(samples, './tmp/cifar10/samples_{}.jpg'.format(frame))
  
  # For calculating inception score
  def get_inception_score(G, ):
      all_samples = []
      for i in xrange(10):
          samples_100 = torch.randn(100, 128)
          if use_cuda:
              samples_100 = samples_100.cuda(gpu)
          samples_100 = autograd.Variable(samples_100, volatile=True)
          all_samples.append(G(samples_100).cpu().data.numpy())
  
      all_samples = np.concatenate(all_samples, axis=0)
      all_samples = np.multiply(np.add(np.multiply(all_samples, 0.5), 0.5), 255).astype('int32')
      all_samples = all_samples.reshape((-1, 3, 32, 32)).transpose(0, 2, 3, 1)
      return lib.inception_score.get_inception_score(list(all_samples))
  
  # Dataset iterator
  train_gen, dev_gen = lib.cifar10.load(BATCH_SIZE, data_dir=DATA_DIR)
  def inf_train_gen():
      while True:
          for images in train_gen():
              # yield images.astype('float32').reshape(BATCH_SIZE, 3, 32, 32).transpose(0, 2, 3, 1)
              yield images
  
  gen = inf_train_gen()
  #preprocess = torchvision.transforms.Compose([
  #                               torchvision.transforms.ToTensor(),
  #                               torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
  #                           ])
  preprocess = torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  
  CRITIC_ITERS = critic_schedule()
  GEN_ITERS = gen_schedule()
  for iteration in range(ITERS):
  
      start_time = time.time()
      ############################
      # (1) Update D network
      ###########################
      for p in netD.parameters():  # reset requires_grad
          p.requires_grad = True  # they are set to False below in netG update
      #for i in range(CRITIC_ITERS):
      for i in range(next(CRITIC_ITERS)):
          _data = next(gen)
          netD.zero_grad()
  
          # train with real
          _data = _data.view((BATCH_SIZE, 3, 32, 32))
          real_data = _data # preprocess(_data)#torch.stack([preprocess(item) for item in _data])
  
          if use_cuda:
              real_data = real_data.cuda(gpu)
          real_data_v = autograd.Variable(real_data, requires_grad=True)
  
          # import torchvision
          # filename = os.path.join("test_train_data", str(iteration) + str(i) + ".jpg")
          # torchvision.utils.save_image(real_data, filename)
  
          D_real = netD(real_data_v)
          D_real = D_real.mean()
          D_real.backward(mone)
  
          # train with fake
          noise = torch.randn(BATCH_SIZE, 128)
          if use_cuda:
              noise = noise.cuda(gpu)
          noisev = autograd.Variable(noise, volatile=True)  # totally freeze netG
          fake = autograd.Variable(netG(noisev).data)
          inputv = fake
          D_fake = netD(inputv)
          D_fake = D_fake.mean()
          D_fake.backward(one)
  
          # train with gradient penalty
          gradient_penalty = calc_gradient_penalty(netD, netG, real_data_v.data, fake.data, real_data_v.grad)
          gradient_penalty.backward()
  
          # print "gradien_penalty: ", gradient_penalty
  
          D_cost = D_fake - D_real + gradient_penalty
          Wasserstein_D = D_real - D_fake
          optimizerD.step()
      ############################
      # (2) Update G network
      ###########################
      for p in netD.parameters():
          p.requires_grad = False  # to avoid computation
      for i in range(next(GEN_ITERS)):
          netG.zero_grad()
  
          noise = torch.randn(BATCH_SIZE, 128)
          if use_cuda:
              noise = noise.cuda(gpu)
          noisev = autograd.Variable(noise)
          fake = netG(noisev)
          G = netD(fake)
          G = G.mean()
          G.backward(mone)
          G_cost = -G
          optimizerG.step()
  
      # Write logs and save samples
      lib.plot.plot('./tmp/cifar10/train disc cost', D_cost.cpu().data.numpy())
      lib.plot.plot('./tmp/cifar10/time', time.time() - start_time)
      lib.plot.plot('./tmp/cifar10/train gen cost', G_cost.cpu().data.numpy())
      lib.plot.plot('./tmp/cifar10/wasserstein distance', Wasserstein_D.cpu().data.numpy())
  
      # Calculate inception score every 1K iters
      if False and iteration % 1000 == 999:
          inception_score = get_inception_score(netG)
          lib.plot.plot('./tmp/cifar10/inception score', inception_score[0])
  
      # Calculate dev loss and generate samples every 100 iters
      if iteration % 100 == 99:
          dev_disc_costs = []
          netD.eval()
          for images in dev_gen():
              images = images.view((BATCH_SIZE, 3, 32, 32))
              imgs = images#preprocess(images)
  
              #imgs = preprocess(images)
              if use_cuda:
                  imgs = imgs.cuda(gpu)
              imgs_v = autograd.Variable(imgs, volatile=True)
  
              D = netD(imgs_v)
              _dev_disc_cost = -D.mean().cpu().data.numpy()
              dev_disc_costs.append(_dev_disc_cost)
          netD.train()
          lib.plot.plot('./tmp/cifar10/dev disc cost', np.mean(dev_disc_costs))
  
          generate_image(iteration, netG)
  
      # Save logs every 100 iters
      if (iteration < 5) or (iteration % 100 == 99):
          lib.plot.flush()
      lib.plot.tick()
  