text: import numpy as np
  import tensorflow as tf
  
  from collections import OrderedDict, defaultdict
  
  from bgan_util import AttributeDict
  
  
  #### Bayesian DCGAN
  
  from dcgan_ops import *
  
  def conv_out_size_same(size, stride):
      return int(math.ceil(float(size) / float(stride)))
  
  
  class BGAN(object):
  
      def __init__(self, x_dim, z_dim, dataset_size, batch_size=64, prior_std=1.0, J=1, M=1, 
                   num_classes=1, alpha=0.01, lr=0.0002, gen_observed=1000,
                   optimizer='adam', wasserstein=False, ml=False):
  
          self.batch_size = batch_size
          self.dataset_size = dataset_size
          self.gen_observed = gen_observed
          self.x_dim = x_dim
          self.z_dim = z_dim
          self.optimizer = optimizer.lower()
  
          self.wasserstein = wasserstein
          if self.wasserstein:
              assert num_classes == 1, "cannot do semi-sup learning with wasserstein ... yet"
              
          # Bayes
          self.prior_std = prior_std
          self.num_gen = J
          self.num_mcmc = M
          self.alpha = alpha
          self.lr = lr
  
          self.ml = ml
          if self.ml:
              assert self.num_gen*self.num_mcmc == 1, "cannot have multiple generators in ml mode"
          
          self.weight_dims = OrderedDict([("g_h0_lin_W", (self.z_dim, 1000)),
                                          ("g_h0_lin_b", (1000,)),
                                          ("g_lin_W", (1000, self.x_dim[0])),
                                          ("g_lin_b", (self.x_dim[0],))])
          
          self.sghmc_noise = {}
          self.noise_std = np.sqrt(2 * self.alpha)
          for name, dim in self.weight_dims.iteritems():
              self.sghmc_noise[name] = tf.contrib.distributions.Normal(mu=0., sigma=self.noise_std*tf.ones(self.weight_dims[name]))
  
          self.K = num_classes # 1 means unsupervised, label == 0 always reserved for fake
  
          self.build_bgan_graph()
  
          if self.K > 1:
              self.build_test_graph()
  
      def _get_optimizer(self, lr):
          if self.optimizer == 'adam':
              return tf.train.AdamOptimizer(learning_rate=lr, beta1=0.5)
          elif self.optimizer == 'sgd':
              return tf.train.MomentumOptimizer(learning_rate=lr, momentum=0.5)
          else:
              raise ValueError("Optimizer must be either 'adam' or 'sgd'")
  
      def build_test_graph(self):
  
          self.test_inputs = tf.placeholder(tf.float32,
                                            [self.batch_size] + self.x_dim, name='real_test_images')
          
          self.lbls = tf.placeholder(tf.float32,
                                     [self.batch_size, self.K], name='real_sup_targets')
  
          self.S, self.S_logits = self.sup_discriminator(self.inputs, self.K)
  
          self.test_D, self.test_D_logits = self.discriminator(self.test_inputs, self.K+1, reuse=True)
          self.test_S, self.test_S_logits = self.sup_discriminator(self.test_inputs, self.K, reuse=True)
  
          self.s_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=self.S_logits,
                                                                               labels=self.lbls))
  
          t_vars = tf.trainable_variables()
          self.sup_vars = [var for var in t_vars if 'sup_' in var.name]
  
          # this is purely supervised
          supervised_lr = 0.05 * self.lr
          s_opt = self._get_optimizer(supervised_lr)
          self.s_optim = s_opt.minimize(self.s_loss, var_list=self.sup_vars)
          s_opt_adam = tf.train.AdamOptimizer(learning_rate=supervised_lr, beta1=0.5)
          self.s_optim_adam = s_opt_adam.minimize(self.s_loss, var_list=self.sup_vars)
  
  
      def build_bgan_graph(self):
      
          self.inputs = tf.placeholder(tf.float32,
                                       [self.batch_size] + self.x_dim, name='real_images')
          
          self.labeled_inputs = tf.placeholder(tf.float32,
                                               [self.batch_size] + self.x_dim, name='real_images_w_labels')
          
          self.labels = tf.placeholder(tf.float32,
                                       [self.batch_size, self.K+1], name='real_targets')
  
  
          self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')
          #self.z_sum = histogram_summary("z", self.z) TODO looks cool
  
          self.gen_param_list = []
          with tf.variable_scope("generator") as scope:
              for gi in xrange(self.num_gen):
                  for m in xrange(self.num_mcmc):
                      gen_params = AttributeDict()
                      for name, shape in self.weight_dims.iteritems():
                          gen_params[name] = tf.get_variable("%s_%04d_%04d" % (name, gi, m),
                                                             shape, initializer=tf.random_normal_initializer(stddev=0.02))
                      self.gen_param_list.append(gen_params)
  
          self.D, self.D_logits = self.discriminator(self.inputs, self.K+1)
          self.Dsup, self.Dsup_logits = self.discriminator(self.labeled_inputs, self.K+1, reuse=True)
  
          if self.K == 1:
              if self.wasserstein:
                  self.d_loss_real = tf.reduce_mean(self.D_logits)
              else:
                  # regular GAN
                  constant_labels = np.zeros((self.batch_size, 2))
                  constant_labels[:, 1] = 1.0
                  self.d_loss_real = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=self.D_logits,
                                                                                            labels=tf.constant(constant_labels)))
          else:
              self.d_loss_sup = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=self.Dsup_logits,
                                                                                       labels=self.labels))    
              self.d_loss_real = -tf.reduce_mean(tf.log((1.0 - self.D[:, 0]) + 1e-8))
                          
  
          self.generation = defaultdict(list)
          for gen_params in self.gen_param_list:
              self.generation["g_prior"].append(self.gen_prior(gen_params))
              self.generation["g_noise"].append(self.gen_noise(gen_params))
              self.generation["generators"].append(self.generator(self.z, gen_params))
              self.generation["gen_samplers"].append(self.sampler(self.z, gen_params))
              D_, D_logits_ = self.discriminator(self.generator(self.z, gen_params), self.K+1, reuse=True)
              self.generation["d_logits"].append(D_logits_)
              self.generation["d_probs"].append(D_)
              
  
          all_d_logits = tf.concat(self.generation["d_logits"], 0)
          if self.wasserstein:
              self.d_loss_fake = -tf.reduce_mean(all_d_logits)
          else:
              constant_labels = np.zeros((self.batch_size*self.num_gen*self.num_mcmc, self.K+1))
              constant_labels[:, 0] = 1.0 # class label indicating it came from generator, aka fake
              self.d_loss_fake = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=all_d_logits,
                                                                                        labels=tf.constant(constant_labels)))
  
          t_vars = tf.trainable_variables()
          self.d_vars = [var for var in t_vars if 'd_' in var.name]
  
          self.d_loss = self.d_loss_real + self.d_loss_fake
          if not self.ml:
              self.d_loss += self.disc_prior() + self.disc_noise()
          if self.K > 1:
              self.d_loss_semi = self.d_loss_sup + self.d_loss_real + self.d_loss_fake
              if not self.ml:
                  self.d_loss_semi += self.disc_prior() + self.disc_noise()
  
          self.g_vars = []
          for gi in xrange(self.num_gen):
              for m in xrange(self.num_mcmc):
                  self.g_vars.append([var for var in t_vars if 'g_' in var.name and "_%04d_%04d" % (gi, m) in var.name])
  
          self.d_learning_rate = tf.placeholder(tf.float32, shape=[])
          d_opt = self._get_optimizer(self.d_learning_rate)
          self.d_optim = d_opt.minimize(self.d_loss, var_list=self.d_vars)
  
          d_opt_adam = tf.train.AdamOptimizer(learning_rate=self.d_learning_rate, beta1=0.5)
          self.d_optim_adam = d_opt_adam.minimize(self.d_loss, var_list=self.d_vars)
          
          clip_d = [w.assign(tf.clip_by_value(w, -0.01, 0.01)) for w in self.d_vars]
          self.clip_d = clip_d
  
          if self.K > 1:
              self.d_semi_learning_rate = tf.placeholder(tf.float32, shape=[])
              d_opt_semi = self._get_optimizer(self.d_semi_learning_rate)
              self.d_optim_semi = d_opt_semi.minimize(self.d_loss_semi, var_list=self.d_vars)
              d_opt_semi_adam = tf.train.AdamOptimizer(learning_rate=self.d_semi_learning_rate, beta1=0.5)
              self.d_optim_semi_adam = d_opt_semi_adam.minimize(self.d_loss_semi, var_list=self.d_vars)
          
          self.g_optims, self.g_optims_adam = [], []
          self.g_learning_rate = tf.placeholder(tf.float32, shape=[])
          for gi in xrange(self.num_gen*self.num_mcmc):
              if self.wasserstein:
                  g_loss = tf.reduce_mean(self.generation["d_logits"][gi])
              else:
                  g_loss = -tf.reduce_mean(tf.log((1.0 - self.generation["d_probs"][gi][:, 0]) + 1e-8))
              if not self.ml:
                  g_loss += self.generation["g_prior"][gi] + self.generation["g_noise"][gi]
              self.generation["g_losses"].append(g_loss)
              g_opt = self._get_optimizer(self.g_learning_rate)
              self.g_optims.append(g_opt.minimize(g_loss, var_list=self.g_vars[gi]))
              g_opt_adam = tf.train.AdamOptimizer(learning_rate=self.g_learning_rate, beta1=0.5)
              self.g_optims_adam.append(g_opt_adam.minimize(g_loss, var_list=self.g_vars[gi]))
  
              
      def discriminator(self, x, K, reuse=False):
          with tf.variable_scope("discriminator") as scope:
              if reuse:
                  scope.reuse_variables()
              h0 = lrelu(linear(x, 1000, 'd_lin_0'))
              h1 = linear(h0, K, 'd_lin_1')
              return tf.nn.softmax(h1), h1
  
      def sup_discriminator(self, x, K, reuse=False):
          
          pass
          
      def generator(self, z, gen_params):
          with tf.variable_scope("generator") as scope:
              h0 = lrelu(linear(z, 1000, 'g_h0_lin',
                                matrix=gen_params.g_h0_lin_W, bias=gen_params.g_h0_lin_b))
              self.x_ = linear(h0, self.x_dim[0], 'g_lin',
                               matrix=gen_params.g_lin_W, bias=gen_params.g_lin_b)
              return self.x_
  
      def sampler(self, z, gen_params):
          with tf.variable_scope("generator") as scope:
              scope.reuse_variables()
              return self.generator(z, gen_params)
  
      def gen_prior(self, gen_params):
          with tf.variable_scope("generator") as scope:
              prior_loss = 0.0
              for var in gen_params.values():
                  nn = tf.divide(var, self.prior_std)
                  prior_loss += tf.reduce_mean(tf.multiply(nn, nn))
                  
          prior_loss /= self.gen_observed
  
          return prior_loss
  
      def gen_noise(self, gen_params): # for SGHMC
          with tf.variable_scope("generator") as scope:
              noise_loss = 0.0
              for name, var in gen_params.iteritems():
                  noise_loss += tf.reduce_sum(var * self.sghmc_noise[name].sample())
          noise_loss /= self.gen_observed
          return noise_loss
  
      def disc_prior(self):
          with tf.variable_scope("discriminator") as scope:
              prior_loss = 0.0
              for var in self.d_vars:
                  nn = tf.divide(var, self.prior_std)
                  prior_loss += tf.reduce_mean(tf.multiply(nn, nn))
                  
          prior_loss /= self.dataset_size
  
          return prior_loss
  
      def disc_noise(self): # for SGHMC
          with tf.variable_scope("discriminator") as scope:
              noise_loss = 0.0
              for var in self.d_vars:
                  noise_ = tf.contrib.distributions.Normal(mu=0., sigma=self.noise_std*tf.ones(var.get_shape()))
                  noise_loss += tf.reduce_sum(var * noise_.sample())
          noise_loss /= self.dataset_size
          return noise_loss
  
  
  class BDCGAN(BGAN):
  
      def __init__(self, x_dim, z_dim, dataset_size, batch_size=64, gf_dim=64, df_dim=64, 
                   prior_std=1.0, J=1, M=1, num_classes=1, eta=2e-4, 
                   alpha=0.01, lr=0.0002, optimizer='adam', wasserstein=False, 
                   ml=False, gen_observed=1000):
  
          assert len(x_dim) == 3, "invalid image dims"
          
          c_dim = x_dim[2]
          self.is_grayscale = (c_dim == 1)
          self.optimizer = optimizer.lower()
          self.dataset_size = dataset_size
          self.batch_size = batch_size
          self.gen_observed = gen_observed
  
          self.x_dim = x_dim
          self.z_dim = z_dim
  
          self.gf_dim = gf_dim
          self.df_dim = df_dim
          self.c_dim = c_dim
          self.lr = lr
  
          self.wasserstein = wasserstein
          
          self.d_bn1 = batch_norm(name='d_bn1')
          self.d_bn2 = batch_norm(name='d_bn2')
          self.d_bn3 = batch_norm(name='d_bn3')
  
          self.sd_bn1 = batch_norm(name='sd_bn1')
          self.sd_bn2 = batch_norm(name='sd_bn2')
          self.sd_bn3 = batch_norm(name='sd_bn3')
  
  
          self.g_bn0 = batch_norm(name='g_bn0')
          self.g_bn1 = batch_norm(name='g_bn1')
          self.g_bn2 = batch_norm(name='g_bn2')
          self.g_bn3 = batch_norm(name='g_bn3')
  
          self.wasserstein = wasserstein
          
          # Bayes
          self.prior_std = prior_std
          self.num_gen = J
          self.num_mcmc = M
          self.eta = eta
          self.alpha = alpha
          # ML
          self.ml = ml
          if self.ml:
              assert self.num_gen == 1, "cannot have >1 generator for ml"
  
          self.output_height = x_dim[0]
          self.output_width = x_dim[1]
  
          s_h, s_w = self.output_height, self.output_width
          s_h2, s_w2 = conv_out_size_same(s_h, 2), conv_out_size_same(s_w, 2)
          s_h4, s_w4 = conv_out_size_same(s_h2, 2), conv_out_size_same(s_w2, 2)
          s_h8, s_w8 = conv_out_size_same(s_h4, 2), conv_out_size_same(s_w4, 2)
          s_h16, s_w16 = conv_out_size_same(s_h8, 2), conv_out_size_same(s_w8, 2)
          
          self.gen_params = AttributeDict()
          self.bgen_params = AttributeDict()
          self.weight_dims = OrderedDict([("g_h0_lin_W", (self.z_dim, self.gf_dim * 8 * s_h16 * s_w16)),
                                          ("g_h0_lin_b", (self.gf_dim * 8 * s_h16 * s_w16,)),
                                          ("g_h1_W", (5, 5, self.gf_dim*4, self.gf_dim*8)),
                                          ("g_h1_b", (self.gf_dim*4,)),
                                          ("g_h2_W", (5, 5, self.gf_dim*2, self.gf_dim*4)),
                                          ("g_h2_b", (self.gf_dim*2,)),
                                          ("g_h3_W", (5, 5, self.gf_dim*1, self.gf_dim*2)),
                                          ("g_h3_b", (self.gf_dim*1,)),
                                          ("g_h4_W", (5, 5, self.c_dim, self.gf_dim*1)),
                                          ("g_h4_b", (self.c_dim,))])
  
          self.sghmc_noise = {}
          self.noise_std = np.sqrt(2 * self.alpha * self.eta)
          for name, dim in self.weight_dims.iteritems():
              self.sghmc_noise[name] = tf.contrib.distributions.Normal(mu=0., sigma=self.noise_std*tf.ones(self.weight_dims[name]))
  
          self.K = num_classes # 1 means unsupervised, label == 0 always reserved for fake
  
          self.build_bgan_graph()
  
          if self.K > 1:
              self.build_test_graph()
                                               
                      
      def discriminator(self, image, K, reuse=False):
          with tf.variable_scope("discriminator") as scope:
              if reuse:
                  scope.reuse_variables()
  
              h0 = lrelu(conv2d(image, self.df_dim, name='d_h0_conv'))
              h1 = lrelu(self.d_bn1(conv2d(h0,
                                           self.df_dim * 2,
                                           name='d_h1_conv')))
              h2 = lrelu(self.d_bn2(conv2d(h1,
                                           self.df_dim * 4,
                                           name='d_h2_conv')))
              h3 = lrelu(self.d_bn3(conv2d(h2,
                                           self.df_dim * 8,
                                           name='d_h3_conv')))
              h4 = linear(tf.reshape(h3, [self.batch_size, -1]), K, 'd_h3_lin')
              return tf.nn.softmax(h4), h4
          
  
      def sup_discriminator(self, image, K, reuse=False):
          with tf.variable_scope("sup_discriminator") as scope:
              if reuse:
                  scope.reuse_variables()
  
              h0 = lrelu(conv2d(image, self.df_dim, name='sup_h0_conv'))
              h1 = lrelu(self.sd_bn1(conv2d(h0,
                                           self.df_dim * 2,
                                           name='sup_h1_conv')))
              h2 = lrelu(self.sd_bn2(conv2d(h1,
                                           self.df_dim * 4,
                                           name='sup_h2_conv')))
              h3 = lrelu(self.sd_bn3(conv2d(h2,
                                           self.df_dim * 8,
                                           name='sup_h3_conv')))
              h4 = linear(tf.reshape(h3, [self.batch_size, -1]), K, 'sup_h3_lin')
              return tf.nn.softmax(h4), h4
  
              
  
      def generator(self, z, gen_params):
          with tf.variable_scope("generator") as scope:
              
              s_h, s_w = self.output_height, self.output_width
              s_h2, s_w2 = conv_out_size_same(s_h, 2), conv_out_size_same(s_w, 2)
              s_h4, s_w4 = conv_out_size_same(s_h2, 2), conv_out_size_same(s_w2, 2)
              s_h8, s_w8 = conv_out_size_same(s_h4, 2), conv_out_size_same(s_w4, 2)
              s_h16, s_w16 = conv_out_size_same(s_h8, 2), conv_out_size_same(s_w8, 2)
  
              # project `z` and reshape
              self.z_, self.h0_w, self.h0_b = linear(z, self.gf_dim * 8 * s_h16 * s_w16, 'g_h0_lin', with_w=True,
                                                     matrix=gen_params.g_h0_lin_W, bias=gen_params.g_h0_lin_b)
  
              self.h0 = tf.reshape(self.z_, [-1, s_h16, s_w16, self.gf_dim * 8])
              h0 = tf.nn.relu(self.g_bn0(self.h0))
  
              self.h1, self.h1_w, self.h1_b = deconv2d(h0,
                                                       [self.batch_size, s_h8, s_w8, self.gf_dim * 4], name='g_h1', with_w=True,
                                                       w=gen_params.g_h1_W, biases=gen_params.g_h1_b)
              h1 = tf.nn.relu(self.g_bn1(self.h1))
  
              h2, self.h2_w, self.h2_b = deconv2d(h1,
                                                  [self.batch_size, s_h4, s_w4, self.gf_dim * 2], name='g_h2', with_w=True,
                                                  w=gen_params.g_h2_W, biases=gen_params.g_h2_b)
              h2 = tf.nn.relu(self.g_bn2(h2))
  
              h3, self.h3_w, self.h3_b = deconv2d(h2,
                                                  [self.batch_size, s_h2, s_w2, self.gf_dim * 1], name='g_h3', with_w=True,
                                                  w=gen_params.g_h3_W, biases=gen_params.g_h3_b)
              h3 = tf.nn.relu(self.g_bn3(h3))
  
              h4, self.h4_w, self.h4_b = deconv2d(h3,
                                                  [self.batch_size, s_h, s_w, self.c_dim], name='g_h4', with_w=True,
                                                  w=gen_params.g_h4_W, biases=gen_params.g_h4_b)
  
              return tf.nn.tanh(h4)
          
  
      def sampler(self, z, gen_params):
          with tf.variable_scope("generator") as scope:
              scope.reuse_variables()
  
              s_h, s_w = self.output_height, self.output_width
              s_h2, s_w2 = conv_out_size_same(s_h, 2), conv_out_size_same(s_w, 2)
              s_h4, s_w4 = conv_out_size_same(s_h2, 2), conv_out_size_same(s_w2, 2)
              s_h8, s_w8 = conv_out_size_same(s_h4, 2), conv_out_size_same(s_w4, 2)
              s_h16, s_w16 = conv_out_size_same(s_h8, 2), conv_out_size_same(s_w8, 2)
  
              # project `z` and reshape
              z_ = linear(z, self.gf_dim * 8 * s_h16 * s_w16, 'g_h0_lin',
                          matrix=gen_params.g_h0_lin_W, bias=gen_params.g_h0_lin_b)
  
              h0 = tf.reshape(z_, [-1, s_h16, s_w16, self.gf_dim * 8])
              h0 = tf.nn.relu(self.g_bn0(h0, train=False))
  
              h1 = deconv2d(h0,
                            [self.batch_size, s_h8, s_w8, self.gf_dim * 4], name='g_h1',
                            w=gen_params.g_h1_W, biases=gen_params.g_h1_b)
              
              h1 = tf.nn.relu(self.g_bn1(h1, train=False))
  
              h2 = deconv2d(h1,
                            [self.batch_size, s_h4, s_w4, self.gf_dim * 2], name='g_h2',
                            w=gen_params.g_h2_W, biases=gen_params.g_h2_b)
              h2 = tf.nn.relu(self.g_bn2(h2, train=False))
  
              h3 = deconv2d(h2,
                            [self.batch_size, s_h2, s_w2, self.gf_dim * 1], name='g_h3',
                            w=gen_params.g_h3_W, biases=gen_params.g_h3_b)
              h3 = tf.nn.relu(self.g_bn3(h3, train=False))
  
              h4 = deconv2d(h3,
                            [self.batch_size, s_h, s_w, self.c_dim], name='g_h4',
                            w=gen_params.g_h4_W, biases=gen_params.g_h4_b)
  
              return tf.nn.tanh(h4)
  
  
  