class_docstr: Define a SketchRNN model.
class_docstr_tok: ["Define","a","SketchRNN","model","."]
docstr: Define model architecture.
docstr_tok: ["Define","model","architecture","."]
code: def build_model(self, hps):
      """Define model architecture."""
      if hps.is_training:
          self.global_step = tf.Variable(
              0, name="global_step", trainable=False
          )
  
      if hps.dec_model == "lstm":
          cell_fn = rnn.LSTMCell
      elif hps.dec_model == "layer_norm":
          cell_fn = rnn.LayerNormLSTMCell
      elif hps.dec_model == "hyper":
          cell_fn = rnn.HyperLSTMCell
      else:
          assert False, "please choose a respectable cell"
  
      if hps.enc_model == "lstm":
          enc_cell_fn = rnn.LSTMCell
      elif hps.enc_model == "layer_norm":
          enc_cell_fn = rnn.LayerNormLSTMCell
      elif hps.enc_model == "hyper":
          enc_cell_fn = rnn.HyperLSTMCell
      else:
          assert False, "please choose a respectable cell"
  
      use_recurrent_dropout = self.hps.use_recurrent_dropout
      use_input_dropout = self.hps.use_input_dropout
      use_output_dropout = self.hps.use_output_dropout
  
      cell = cell_fn(
          hps.dec_rnn_size,
          use_recurrent_dropout=use_recurrent_dropout,
          dropout_keep_prob=self.hps.recurrent_dropout_prob,
      )
  
      if hps.conditional:  # vae mode:
          if hps.enc_model == "hyper":
              self.enc_cell_fw = enc_cell_fn(
                  hps.enc_rnn_size,
                  use_recurrent_dropout=use_recurrent_dropout,
                  dropout_keep_prob=self.hps.recurrent_dropout_prob,
              )
              self.enc_cell_bw = enc_cell_fn(
                  hps.enc_rnn_size,
                  use_recurrent_dropout=use_recurrent_dropout,
                  dropout_keep_prob=self.hps.recurrent_dropout_prob,
              )
          else:
              self.enc_cell_fw = enc_cell_fn(
                  hps.enc_rnn_size,
                  use_recurrent_dropout=use_recurrent_dropout,
                  dropout_keep_prob=self.hps.recurrent_dropout_prob,
              )
              self.enc_cell_bw = enc_cell_fn(
                  hps.enc_rnn_size,
                  use_recurrent_dropout=use_recurrent_dropout,
                  dropout_keep_prob=self.hps.recurrent_dropout_prob,
              )
  
      # dropout:
      tf.logging.info("Input dropout mode = %s.", use_input_dropout)
      tf.logging.info("Output dropout mode = %s.", use_output_dropout)
      tf.logging.info(
          "Recurrent dropout mode = %s.", use_recurrent_dropout
      )
      if use_input_dropout:
          tf.logging.info(
              "Dropout to input w/ keep_prob = %4.4f.",
              self.hps.input_dropout_prob,
          )
          cell = tf.nn.rnn_cell.DropoutWrapper(
              cell, input_keep_prob=self.hps.input_dropout_prob
          )
      if use_output_dropout:
          tf.logging.info(
              "Dropout to output w/ keep_prob = %4.4f.",
              self.hps.output_dropout_prob,
          )
          cell = tf.nn.rnn_cell.DropoutWrapper(
              cell, output_keep_prob=self.hps.output_dropout_prob
          )
      self.cell = cell
  
      self.sequence_lengths = tf.placeholder(
          dtype=tf.int32, shape=[self.hps.batch_size]
      )
      self.input_data = tf.placeholder(
          dtype=tf.float32,
          shape=[self.hps.batch_size, self.hps.max_seq_len + 1, 5],
      )
  
      # The target/expected vectors of strokes
      self.output_x = self.input_data[
          :, 1 : self.hps.max_seq_len + 1, :
      ]
      # vectors of strokes to be fed to decoder (same as above, but lagged behind
      # one step to include initial dummy value of (0, 0, 1, 0, 0))
      self.input_x = self.input_data[:, : self.hps.max_seq_len, :]
  
      # either do vae-bit and get z, or do unconditional, decoder-only
      if hps.conditional:  # vae mode:
          self.mean, self.presig = self.encoder(
              self.output_x, self.sequence_lengths
          )
          self.sigma = tf.exp(
              self.presig / 2.0
          )  # sigma > 0. div 2.0 -> sqrt.
          eps = tf.random_normal(
              (self.hps.batch_size, self.hps.z_size),
              0.0,
              1.0,
              dtype=tf.float32,
          )
          self.batch_z = self.mean + tf.multiply(self.sigma, eps)
          # KL cost
          self.kl_cost = -0.5 * tf.reduce_mean(
              (
                  1
                  + self.presig
                  - tf.square(self.mean)
                  - tf.exp(self.presig)
              )
          )
          self.kl_cost = tf.maximum(self.kl_cost, self.hps.kl_tolerance)
          pre_tile_y = tf.reshape(
              self.batch_z, [self.hps.batch_size, 1, self.hps.z_size]
          )
          overlay_x = tf.tile(pre_tile_y, [1, self.hps.max_seq_len, 1])
          actual_input_x = tf.concat([self.input_x, overlay_x], 2)
          self.initial_state = tf.nn.tanh(
              rnn.super_linear(
                  self.batch_z,
                  cell.state_size,
                  init_w="gaussian",
                  weight_start=0.001,
                  input_size=self.hps.z_size,
              )
          )
      else:  # unconditional, decoder-only generation
          self.batch_z = tf.zeros(
              (self.hps.batch_size, self.hps.z_size), dtype=tf.float32
          )
          self.kl_cost = tf.zeros([], dtype=tf.float32)
          actual_input_x = self.input_x
          self.initial_state = cell.zero_state(
              batch_size=hps.batch_size, dtype=tf.float32
          )
  
      self.num_mixture = hps.num_mixture
  
      # TODO(deck): Better understand this comment.
      # Number of outputs is 3 (one logit per pen state) plus 6 per mixture
      # component: mean_x, stdev_x, mean_y, stdev_y, correlation_xy, and the
      # mixture weight/probability (Pi_k)
      n_out = 3 + self.num_mixture * 6
  
      with tf.variable_scope("RNN"):
          output_w = tf.get_variable(
              "output_w", [self.hps.dec_rnn_size, n_out]
          )
          output_b = tf.get_variable("output_b", [n_out])
  
      # decoder module of sketch-rnn is below
      output, last_state = tf.nn.dynamic_rnn(
          cell,
          actual_input_x,
          initial_state=self.initial_state,
          time_major=False,
          swap_memory=True,
          dtype=tf.float32,
          scope="RNN",
      )
  
      output = tf.reshape(output, [-1, hps.dec_rnn_size])
      output = tf.nn.xw_plus_b(output, output_w, output_b)
      self.final_state = last_state
  
      # NB: the below are inner functions, not methods of Model
      def tf_2d_normal(x1, x2, mu1, mu2, s1, s2, rho):
          """Returns result of eq # 24 of http://arxiv.org/abs/1308.0850."""
          norm1 = tf.subtract(x1, mu1)
          norm2 = tf.subtract(x2, mu2)
          s1s2 = tf.multiply(s1, s2)
          # eq 25
          z = (
              tf.square(tf.div(norm1, s1))
              + tf.square(tf.div(norm2, s2))
              - 2
              * tf.div(
                  tf.multiply(rho, tf.multiply(norm1, norm2)), s1s2
              )
          )
          neg_rho = 1 - tf.square(rho)
          result = tf.exp(tf.div(-z, 2 * neg_rho))
          denom = 2 * np.pi * tf.multiply(s1s2, tf.sqrt(neg_rho))
          result = tf.div(result, denom)
          return result
  
      def get_lossfunc(
          z_pi,
          z_mu1,
          z_mu2,
          z_sigma1,
          z_sigma2,
          z_corr,
          z_pen_logits,
          x1_data,
          x2_data,
          pen_data,
      ):
          """Returns a loss fn based on eq #26 of http://arxiv.org/abs/1308.0850."""
          # This represents the L_R only (i.e. does not include the KL loss term).
  
          result0 = tf_2d_normal(
              x1_data, x2_data, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr
          )
          epsilon = 1e-6
          # result1 is the loss wrt pen offset (L_s in equation 9 of
          # https://arxiv.org/pdf/1704.03477.pdf)
          result1 = tf.multiply(result0, z_pi)
          result1 = tf.reduce_sum(result1, 1, keep_dims=True)
          result1 = -tf.log(result1 + epsilon)  # avoid log(0)
  
          fs = 1.0 - pen_data[:, 2]  # use training data for this
          fs = tf.reshape(fs, [-1, 1])
          # Zero out loss terms beyond N_s, the last actual stroke
          result1 = tf.multiply(result1, fs)
  
          # result2: loss wrt pen state, (L_p in equation 9)
          result2 = tf.nn.softmax_cross_entropy_with_logits(
              labels=pen_data, logits=z_pen_logits
          )
          result2 = tf.reshape(result2, [-1, 1])
          if not self.hps.is_training:  # eval mode, mask eos columns
              result2 = tf.multiply(result2, fs)
  
          result = result1 + result2
          return result
  
      # below is where we need to do MDN (Mixture Density Network) splitting of
      # distribution params
      def get_mixture_coef(output):
          """Returns the tf slices containing mdn dist params."""
          # This uses eqns 18 -> 23 of http://arxiv.org/abs/1308.0850.
          z = output
          z_pen_logits = z[:, 0:3]  # pen states
          z_pi, z_mu1, z_mu2, z_sigma1, z_sigma2, z_corr = tf.split(
              z[:, 3:], 6, 1
          )
  
          # process output z's into MDN parameters
  
          # softmax all the pi's and pen states:
          z_pi = tf.nn.softmax(z_pi)
          z_pen = tf.nn.softmax(z_pen_logits)
  
          # exponentiate the sigmas and also make corr between -1 and 1.
          z_sigma1 = tf.exp(z_sigma1)
          z_sigma2 = tf.exp(z_sigma2)
          z_corr = tf.tanh(z_corr)
  
          r = [
              z_pi,
              z_mu1,
              z_mu2,
              z_sigma1,
              z_sigma2,
              z_corr,
              z_pen,
              z_pen_logits,
          ]
          return r
  
      out = get_mixture_coef(output)
      [
          o_pi,
          o_mu1,
          o_mu2,
          o_sigma1,
          o_sigma2,
          o_corr,
          o_pen,
          o_pen_logits,
      ] = out
  
      self.pi = o_pi
      self.mu1 = o_mu1
      self.mu2 = o_mu2
      self.sigma1 = o_sigma1
      self.sigma2 = o_sigma2
      self.corr = o_corr
      self.pen_logits = o_pen_logits
      # pen state probabilities (result of applying softmax to self.pen_logits)
      self.pen = o_pen
  
      # reshape target data so that it is compatible with prediction shape
      target = tf.reshape(self.output_x, [-1, 5])
      [x1_data, x2_data, eos_data, eoc_data, cont_data] = tf.split(
          target, 5, 1
      )
      pen_data = tf.concat([eos_data, eoc_data, cont_data], 1)
  
      lossfunc = get_lossfunc(
          o_pi,
          o_mu1,
          o_mu2,
          o_sigma1,
          o_sigma2,
          o_corr,
          o_pen_logits,
          x1_data,
          x2_data,
          pen_data,
      )
  
      self.r_cost = tf.reduce_mean(lossfunc)
  
      if self.hps.is_training:
          self.lr = tf.Variable(self.hps.learning_rate, trainable=False)
          optimizer = tf.train.AdamOptimizer(self.lr)
  
          self.kl_weight = tf.Variable(
              self.hps.kl_weight_start, trainable=False
          )
          self.cost = self.r_cost + self.kl_cost * self.kl_weight
  
          gvs = optimizer.compute_gradients(self.cost)
          g = self.hps.grad_clip
          capped_gvs = [
              (tf.clip_by_value(grad, -g, g), var) for grad, var in gvs
          ]
          self.train_op = optimizer.apply_gradients(
              capped_gvs,
              global_step=self.global_step,
              name="train_step",
          )
  
file_path: magenta/magenta/models/sketch_rnn/model.py
imports: import random
  from magenta.contrib import training
  from magenta.models.sketch_rnn import rnn
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling the Model obj's underlying member method code with example usage: obj.build_model(hps) and returns: result, result, r