rnn_gan.py

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""

The hyperparameters used in the model:
- learning_rate - the initial value of the learning rate
- max_grad_norm - the maximum permissible norm of the gradient
- num_layers - the number of LSTM layers
- songlength - the number of unrolled steps of LSTM
- hidden_size - the number of LSTM units
- epochs_before_decay - the number of epochs trained with the initial learning rate
- max_epoch - the total number of epochs for training
- keep_prob - the probability of keeping weights in the dropout layer
- lr_decay - the decay of the learning rate for each epoch after "epochs_before_decay"
- batch_size - the batch size

The hyperparameters that could be used in the model:
- init_scale - the initial scale of the weights

To run:

$ python rnn_gan.py --model small|medium|large --datadir simple-examples/data/ --traindir dir-for-checkpoints-and-plots --select_validation_percentage 0-40 --select_test_percentage 0-40

"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time, datetime, os, sys
import cPickle as pkl
from subprocess import call, Popen

import numpy as np
import tensorflow as tf
from tensorflow.python.client import timeline


import music_data_utils
from midi_statistics import get_all_stats

flags = tf.flags
logging = tf.logging

flags.DEFINE_string("datadir", None, "Directory to save and load midi music files.")
flags.DEFINE_string("traindir", None, "Directory to save checkpoints and gnuplot files.")
flags.DEFINE_integer("epochs_per_checkpoint", 2,
                     "How many training epochs to do per checkpoint.")
flags.DEFINE_boolean("log_device_placement", False,           #
                   "Outputs info on device placement.")
flags.DEFINE_string("call_after", None, "Call this command after exit.")
flags.DEFINE_integer("exit_after", 1440,
                     "exit after this many minutes")
flags.DEFINE_integer("select_validation_percentage", None,
                     "Select random percentage of data as validation set.")
flags.DEFINE_integer("select_test_percentage", None,
                     "Select random percentage of data as test set.")
flags.DEFINE_boolean("sample", False,
                     "Sample output from the model. Assume training was already done. Save sample output to file.")
flags.DEFINE_integer("works_per_composer", None,
                     "Limit number of works per composer that is loaded.")
flags.DEFINE_boolean("disable_feed_previous", False,
                     "Feed output from previous cell to the input of the next. In the generator.")
flags.DEFINE_float("init_scale", 0.05,                # .1, .04
                   "the initial scale of the weights")
flags.DEFINE_float("learning_rate", 0.1,              # .05,.1,.9 
                   "Learning rate")
flags.DEFINE_float("d_lr_factor", 0.5,                # .5
                   "Learning rate decay")
flags.DEFINE_float("max_grad_norm", 5.0,              # 5.0, 10.0
                   "the maximum permissible norm of the gradient")
flags.DEFINE_float("keep_prob", 0.5,                  # 1.0, .35
                   "Keep probability. 1.0 disables dropout.")
flags.DEFINE_float("lr_decay", 1.0,                   # 1.0
                   "Learning rate decay after each epoch after epochs_before_decay")
flags.DEFINE_integer("num_layers_g", 2,                 # 2
                   "Number of stacked recurrent cells in G.")
flags.DEFINE_integer("num_layers_d", 2,                 # 2
                   "Number of stacked recurrent cells in D.")
flags.DEFINE_integer("songlength", 100,               # 200, 500
                   "Limit song inputs to this number of events.")
flags.DEFINE_integer("meta_layer_size", 200,          # 300, 600
                   "Size of hidden layer for meta information module.")
flags.DEFINE_integer("hidden_size_g", 350,              # 200, 1500
                   "Hidden size for recurrent part of G.")
flags.DEFINE_integer("hidden_size_d", 350,              # 200, 1500
                   "Hidden size for recurrent part of D. Default: same as for G.")
flags.DEFINE_integer("epochs_before_decay", 60,       # 40, 140
                   "Number of epochs before starting to decay.")
flags.DEFINE_integer("max_epoch", 500,                # 500, 500
                   "Number of epochs before stopping training.")
flags.DEFINE_integer("batch_size", 20,                # 10, 20
                   "Batch size.")
flags.DEFINE_integer("biscale_slow_layer_ticks", 8,   # 8
                   "Biscale slow layer ticks. Not implemented yet.")
flags.DEFINE_boolean("multiscale", False,             #
                   "Multiscale RNN. Not implemented.")
flags.DEFINE_integer("pretraining_epochs", 6,        # 20, 40
                   "Number of epochs to run lang-model style pretraining.")
flags.DEFINE_boolean("pretraining_d", False,          #
                   "Train D during pretraining.")
flags.DEFINE_boolean("initialize_d", False,           #
                   "Initialize variables for D, no matter if there are trained versions in checkpoint.")
flags.DEFINE_boolean("ignore_saved_args", False,      #
                   "Tells the program to ignore saved arguments, and instead use the ones provided as CLI arguments.")
flags.DEFINE_boolean("pace_events", False,            #
                   "When parsing input data, insert one dummy event at each quarter note if there is no tone.")
flags.DEFINE_boolean("minibatch_d", False,            #
                   "Adding kernel features for minibatch diversity.")
flags.DEFINE_boolean("unidirectional_d", False,        #
                   "Unidirectional RNN instead of bidirectional RNN for D.")
flags.DEFINE_boolean("profiling", False,              #
                   "Profiling. Writing a timeline.json file in plots dir.")
flags.DEFINE_boolean("float16", False,                #
                   "Use floa16 data type. Otherwise, use float32.")

flags.DEFINE_boolean("adam", False,                   #
                   "Use Adam optimizer.")
flags.DEFINE_boolean("feature_matching", False,       #
                   "Feature matching objective for G.")
flags.DEFINE_boolean("disable_l2_regularizer", False,       #
                   "L2 regularization on weights.")
flags.DEFINE_float("reg_scale", 1.0,       #
                   "L2 regularization scale.")
flags.DEFINE_boolean("synthetic_chords", False,       #
                   "Train on synthetically generated chords (three tones per event).")
flags.DEFINE_integer("tones_per_cell", 1,             # 2,3
                   "Maximum number of tones to output per RNN cell.")
flags.DEFINE_string("composer", None, "Specify exactly one composer, and train model only on this.")
flags.DEFINE_boolean("generate_meta", False, "Generate the composer and genre as part of output.")
flags.DEFINE_float("random_input_scale", 1.0,       #
                   "Scale of random inputs (1.0=same size as generated features).")
flags.DEFINE_boolean("end_classification", False, "Classify only in ends of D. Otherwise, does classification at every timestep and mean reduce.")

FLAGS = flags.FLAGS

model_layout_flags = ['num_layers_g', 'num_layers_d', 'meta_layer_size', 'hidden_size_g', 'hidden_size_d', 'biscale_slow_layer_ticks', 'multiscale', 'multiscale', 'disable_feed_previous', 'pace_events', 'minibatch_d', 'unidirectional_d', 'feature_matching', 'composer']


def restore_flags(save_if_none_found=True):
  if FLAGS.traindir:
    saved_args_dir = os.path.join(FLAGS.traindir, 'saved_args')
    if save_if_none_found:
      try: os.makedirs(saved_args_dir)
      except: pass
    for arg in FLAGS.__flags:
      if arg not in model_layout_flags:
        continue
      if FLAGS.ignore_saved_args and os.path.exists(os.path.join(saved_args_dir, arg+'.pkl')):
        print('{:%Y-%m-%d %H:%M:%S}: saved_args: Found {} setting from saved state, but using CLI args ({}) and saving (--ignore_saved_args).'.format(datetime.datetime.today(), arg, getattr(FLAGS, arg)))
      elif os.path.exists(os.path.join(saved_args_dir, arg+'.pkl')):
        with open(os.path.join(saved_args_dir, arg+'.pkl'), 'r') as f:
          setattr(FLAGS, arg, pkl.load(f))
          print('{:%Y-%m-%d %H:%M:%S}: saved_args: {} from saved state ({}), ignoring CLI args.'.format(datetime.datetime.today(), arg, getattr(FLAGS, arg)))
      elif save_if_none_found:
        print('{:%Y-%m-%d %H:%M:%S}: saved_args: Found no {} setting from saved state, using CLI args ({}) and saving.'.format(datetime.datetime.today(), arg, getattr(FLAGS, arg)))
        with open(os.path.join(saved_args_dir, arg+'.pkl'), 'w') as f:
          pkl.dump(getattr(FLAGS, arg), f)
      else:
        print('{:%Y-%m-%d %H:%M:%S}: saved_args: Found no {} setting from saved state, using CLI args ({}) but not saving.'.format(datetime.datetime.today(), arg, getattr(FLAGS, arg)))



def data_type():
  return tf.float16 if FLAGS.float16 else tf.float32
  #return tf.float16

def my_reduce_mean(what_to_take_mean_over):
  return tf.reshape(what_to_take_mean_over, shape=[-1])[0]
  denom = 1.0
  #print(what_to_take_mean_over.get_shape())
  for d in what_to_take_mean_over.get_shape():
    #print(d)
    if type(d) == tf.Dimension:
      denom = denom*d.value
    else:
      denom = denom*d
  return tf.reduce_sum(what_to_take_mean_over)/denom

def linear(inp, output_dim, scope=None, stddev=1.0, reuse_scope=False):
  norm = tf.random_normal_initializer(stddev=stddev, dtype=data_type())
  const = tf.constant_initializer(0.0, dtype=data_type())
  with tf.variable_scope(scope or 'linear') as scope:
    scope.set_regularizer(tf.contrib.layers.l2_regularizer(scale=FLAGS.reg_scale))
    if reuse_scope:
      scope.reuse_variables()
    #print('inp.get_shape(): {}'.format(inp.get_shape()))
    w = tf.get_variable('w', [inp.get_shape()[1], output_dim], initializer=norm, dtype=data_type())
    b = tf.get_variable('b', [output_dim], initializer=const, dtype=data_type())
  return tf.matmul(inp, w) + b

def minibatch(inp, num_kernels=25, kernel_dim=10, scope=None, msg='', reuse_scope=False):
  """
   Borrowed from http://blog.aylien.com/introduction-generative-adversarial-networks-code-tensorflow/
  """
  with tf.variable_scope(scope or 'minibatch_d') as scope:
    scope.set_regularizer(tf.contrib.layers.l2_regularizer(scale=FLAGS.reg_scale))
    if reuse_scope:
      scope.reuse_variables()
  
    inp = tf.Print(inp, [inp],
            '{} inp = '.format(msg), summarize=20, first_n=20)
    x = tf.sigmoid(linear(inp, num_kernels * kernel_dim, scope))
    activation = tf.reshape(x, (-1, num_kernels, kernel_dim))
    activation = tf.Print(activation, [activation],
            '{} activation = '.format(msg), summarize=20, first_n=20)
    diffs = tf.expand_dims(activation, 3) - \
                tf.expand_dims(tf.transpose(activation, [1, 2, 0]), 0)
    diffs = tf.Print(diffs, [diffs],
            '{} diffs = '.format(msg), summarize=20, first_n=20)
    abs_diffs = tf.reduce_sum(tf.abs(diffs), 2)
    abs_diffs = tf.Print(abs_diffs, [abs_diffs],
            '{} abs_diffs = '.format(msg), summarize=20, first_n=20)
    minibatch_features = tf.reduce_sum(tf.exp(-abs_diffs), 2)
    minibatch_features = tf.Print(minibatch_features, [tf.reduce_min(minibatch_features), tf.reduce_max(minibatch_features)],
            '{} minibatch_features (min,max) = '.format(msg), summarize=20, first_n=20)
  return tf.concat(1, [inp, minibatch_features])

class RNNGAN(object):
  """The RNNGAN model."""

  def __init__(self, is_training, num_song_features=None, num_meta_features=None):
    self.batch_size = batch_size = FLAGS.batch_size
    self.songlength = songlength = FLAGS.songlength
    #self.global_step            = tf.Variable(0, trainable=False)

    print('songlength: {}'.format(self.songlength))
    self._input_songdata = tf.placeholder(shape=[batch_size, songlength, num_song_features], dtype=data_type())
    self._input_metadata = tf.placeholder(shape=[batch_size, num_meta_features], dtype=data_type())

    songdata_inputs = [tf.squeeze(input_, [1])
              for input_ in tf.split(1, songlength, self._input_songdata)]
    
    with tf.variable_scope('G') as scope:
      scope.set_regularizer(tf.contrib.layers.l2_regularizer(scale=FLAGS.reg_scale))
      lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(FLAGS.hidden_size_g, forget_bias=1.0, state_is_tuple=True)
      if is_training and FLAGS.keep_prob < 1:
        lstm_cell = tf.nn.rnn_cell.DropoutWrapper(
            lstm_cell, output_keep_prob=FLAGS.keep_prob)
      cell = tf.nn.rnn_cell.MultiRNNCell([lstm_cell] * FLAGS.num_layers_g, state_is_tuple=True)

      self._initial_state = cell.zero_state(batch_size, data_type())

      # TODO: (possibly temporarily) disabling meta info
      if FLAGS.generate_meta:
        metainputs = tf.random_uniform(shape=[batch_size, int(FLAGS.random_input_scale*num_meta_features)], minval=0.0, maxval=1.0)
        meta_g = tf.nn.relu(linear(metainputs, FLAGS.meta_layer_size, scope='meta_layer', reuse_scope=False))
        meta_softmax_w = tf.get_variable("meta_softmax_w", [FLAGS.meta_layer_size, num_meta_features])
        meta_softmax_b = tf.get_variable("meta_softmax_b", [num_meta_features])
        meta_logits = tf.nn.xw_plus_b(meta_g, meta_softmax_w, meta_softmax_b)
        meta_probs = tf.nn.softmax(meta_logits)

      random_rnninputs = tf.random_uniform(shape=[batch_size, songlength, int(FLAGS.random_input_scale*num_song_features)], minval=0.0, maxval=1.0, dtype=data_type())

      # Make list of tensors. One per step in recurrence.
      # Each tensor is batchsize*numfeatures.
      random_rnninputs = [tf.squeeze(input_, [1]) for input_ in tf.split(1, songlength, random_rnninputs)]
      
      # REAL GENERATOR:
      state = self._initial_state
      # as we feed the output as the input to the next, we 'invent' the initial 'output'.
      generated_point = tf.random_uniform(shape=[batch_size, num_song_features], minval=0.0, maxval=1.0, dtype=data_type())
      outputs = []
      self._generated_features = []
      for i,input_ in enumerate(random_rnninputs):
        if i > 0: scope.reuse_variables()
        concat_values = [input_]
        if not FLAGS.disable_feed_previous:
          concat_values.append(generated_point)
        if FLAGS.generate_meta:
          concat_values.append(meta_probs)
        if len(concat_values):
          input_ = tf.concat(concat_dim=1, values=concat_values)
        input_ = tf.nn.relu(linear(input_, FLAGS.hidden_size_g,
                            scope='input_layer', reuse_scope=(i!=0)))
        output, state = cell(input_, state)
        outputs.append(output)
        #generated_point = tf.nn.relu(linear(output, num_song_features, scope='output_layer', reuse_scope=(i!=0)))
        generated_point = linear(output, num_song_features, scope='output_layer', reuse_scope=(i!=0))
        self._generated_features.append(generated_point)
      
      
      # PRETRAINING GENERATOR, will feed inputs, not generated outputs:
      scope.reuse_variables()
      # as we feed the output as the input to the next, we 'invent' the initial 'output'.
      prev_target = tf.random_uniform(shape=[batch_size, num_song_features], minval=0.0, maxval=1.0, dtype=data_type())
      outputs = []
      self._generated_features_pretraining = []
      for i,input_ in enumerate(random_rnninputs):
        concat_values = [input_]
        if not FLAGS.disable_feed_previous:
          concat_values.append(prev_target)
        if FLAGS.generate_meta:
          concat_values.append(self._input_metadata)
        if len(concat_values):
          input_ = tf.concat(concat_dim=1, values=concat_values)
        input_ = tf.nn.relu(linear(input_, FLAGS.hidden_size_g, scope='input_layer', reuse_scope=(i!=0)))
        output, state = cell(input_, state)
        outputs.append(output)
        #generated_point = tf.nn.relu(linear(output, num_song_features, scope='output_layer', reuse_scope=(i!=0)))
        generated_point = linear(output, num_song_features, scope='output_layer', reuse_scope=(i!=0))
        self._generated_features_pretraining.append(generated_point)
        prev_target = songdata_inputs[i]
      
      #outputs, state = tf.nn.rnn(cell, transformed, initial_state=self._initial_state)

      #self._generated_features = [tf.nn.relu(linear(output, num_song_features, scope='output_layer', reuse_scope=(i!=0))) for i,output in enumerate(outputs)]

    self._final_state = state

    # These are used both for pretraining and for D/G training further down.
    self._lr = tf.Variable(FLAGS.learning_rate, trainable=False, dtype=data_type())
    self.g_params = [v for v in tf.trainable_variables() if v.name.startswith('model/G/')]
    if FLAGS.adam:
      g_optimizer = tf.train.AdamOptimizer(self._lr)
    else:
      g_optimizer = tf.train.GradientDescentOptimizer(self._lr)
   
    reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
    reg_constant = 0.1  # Choose an appropriate one.
    reg_loss = reg_constant * sum(reg_losses)
    reg_loss = tf.Print(reg_loss, reg_losses,
                  'reg_losses = ', summarize=20, first_n=20)
    #if not FLAGS.disable_l2_regularizer:
    #  print('L2 regularization. Reg losses: {}'.format([v.name for v in reg_losses]))
   
    # ---BEGIN, PRETRAINING. ---
    
    print(tf.transpose(tf.pack(self._generated_features_pretraining), perm=[1, 0, 2]).get_shape())
    print(self._input_songdata.get_shape())
    self.rnn_pretraining_loss = tf.reduce_mean(tf.squared_difference(x=tf.transpose(tf.pack(self._generated_features_pretraining), perm=[1, 0, 2]), y=self._input_songdata))
    if not FLAGS.disable_l2_regularizer:
      self.rnn_pretraining_loss = self.rnn_pretraining_loss+reg_loss
    
    
    pretraining_grads, _ = tf.clip_by_global_norm(tf.gradients(self.rnn_pretraining_loss, self.g_params), FLAGS.max_grad_norm)
    self.opt_pretraining = g_optimizer.apply_gradients(zip(pretraining_grads, self.g_params))

    # ---END, PRETRAINING---

    # The discriminator tries to tell the difference between samples from the
    # true data distribution (self.x) and the generated samples (self.z).
    #
    # Here we create two copies of the discriminator network (that share parameters),
    # as you cannot use the same network with different inputs in TensorFlow.
    with tf.variable_scope('D') as scope:
      scope.set_regularizer(tf.contrib.layers.l2_regularizer(scale=FLAGS.reg_scale))
      # Make list of tensors. One per step in recurrence.
      # Each tensor is batchsize*numfeatures.
      # TODO: (possibly temporarily) disabling meta info
      print('self._input_songdata shape {}'.format(self._input_songdata.get_shape()))
      print('generated data shape {}'.format(self._generated_features[0].get_shape()))
      # TODO: (possibly temporarily) disabling meta info
      if FLAGS.generate_meta:
        songdata_inputs = [tf.concat(1, [self._input_metadata, songdata_input]) for songdata_input in songdata_inputs]
      #print('metadata inputs shape {}'.format(self._input_metadata.get_shape()))
      #print('generated metadata shape {}'.format(meta_probs.get_shape()))
      self.real_d,self.real_d_features = self.discriminator(songdata_inputs, is_training, msg='real')
      scope.reuse_variables()
      # TODO: (possibly temporarily) disabling meta info
      if FLAGS.generate_meta:
        generated_data = [tf.concat(1, [meta_probs, songdata_input]) for songdata_input in self._generated_features]
      else:
        generated_data = self._generated_features
      if songdata_inputs[0].get_shape() != generated_data[0].get_shape():
        print('songdata_inputs shape {} != generated data shape {}'.format(songdata_inputs[0].get_shape(), generated_data[0].get_shape()))
      self.generated_d,self.generated_d_features = self.discriminator(generated_data, is_training, msg='generated')

    # Define the loss for discriminator and generator networks (see the original
    # paper for details), and create optimizers for both
    self.d_loss = tf.reduce_mean(-tf.log(tf.clip_by_value(self.real_d, 1e-1000000, 1.0)) \
                                 -tf.log(1 - tf.clip_by_value(self.generated_d, 0.0, 1.0-1e-1000000)))
    self.g_loss_feature_matching = tf.reduce_sum(tf.squared_difference(self.real_d_features, self.generated_d_features))
    self.g_loss = tf.reduce_mean(-tf.log(tf.clip_by_value(self.generated_d, 1e-1000000, 1.0)))

    if not FLAGS.disable_l2_regularizer:
      self.d_loss = self.d_loss+reg_loss
      self.g_loss_feature_matching = self.g_loss_feature_matching+reg_loss
      self.g_loss = self.g_loss+reg_loss
    self.d_params = [v for v in tf.trainable_variables() if v.name.startswith('model/D/')]

    if not is_training:
      return

    d_optimizer = tf.train.GradientDescentOptimizer(self._lr*FLAGS.d_lr_factor)
    d_grads, _ = tf.clip_by_global_norm(tf.gradients(self.d_loss, self.d_params),
                                        FLAGS.max_grad_norm)
    self.opt_d = d_optimizer.apply_gradients(zip(d_grads, self.d_params))
    if FLAGS.feature_matching:
      g_grads, _ = tf.clip_by_global_norm(tf.gradients(self.g_loss_feature_matching,
                                                       self.g_params),
                                        FLAGS.max_grad_norm)
    else:
      g_grads, _ = tf.clip_by_global_norm(tf.gradients(self.g_loss, self.g_params),
                                        FLAGS.max_grad_norm)
    self.opt_g = g_optimizer.apply_gradients(zip(g_grads, self.g_params))

    self._new_lr = tf.placeholder(shape=[], name="new_learning_rate", dtype=data_type())
    self._lr_update = tf.assign(self._lr, self._new_lr)

  def discriminator(self, inputs, is_training, msg=''):
    # RNN discriminator:
    #for i in xrange(len(inputs)):
    #  print('shape inputs[{}] {}'.format(i, inputs[i].get_shape()))
    #inputs[0] = tf.Print(inputs[0], [inputs[0]],
    #        '{} inputs[0] = '.format(msg), summarize=20, first_n=20)
    if is_training and FLAGS.keep_prob < 1:
      inputs = [tf.nn.dropout(inp, FLAGS.keep_prob) for inp in inputs]
    lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(FLAGS.hidden_size_d, forget_bias=1.0, state_is_tuple=True)
    if is_training and FLAGS.keep_prob < 1:
      lstm_cell = tf.nn.rnn_cell.DropoutWrapper(
          lstm_cell, output_keep_prob=FLAGS.keep_prob)
    cell_fw = tf.nn.rnn_cell.MultiRNNCell([lstm_cell] * FLAGS.num_layers_d, state_is_tuple=True)
    self._initial_state_fw = cell_fw.zero_state(self.batch_size, data_type())
    if not FLAGS.unidirectional_d:
      lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(FLAGS.hidden_size_g, forget_bias=1.0, state_is_tuple=True)
      if is_training and FLAGS.keep_prob < 1:
        lstm_cell = tf.nn.rnn_cell.DropoutWrapper(
            lstm_cell, output_keep_prob=FLAGS.keep_prob)
      cell_bw = tf.nn.rnn_cell.MultiRNNCell([lstm_cell] * FLAGS.num_layers_d, state_is_tuple=True)
      self._initial_state_bw = cell_bw.zero_state(self.batch_size, data_type())
      
      outputs, state_fw, state_bw = tf.nn.bidirectional_rnn(cell_fw, cell_bw, inputs, initial_state_fw=self._initial_state_fw, initial_state_bw=self._initial_state_bw)
      #outputs[0] = tf.Print(outputs[0], [outputs[0]],
      #        '{} outputs[0] = '.format(msg), summarize=20, first_n=20)
      #state = tf.concat(state_fw, state_bw)
      #endoutput = tf.concat(concat_dim=1, values=[outputs[0],outputs[-1]])
    else:
      outputs, state = tf.nn.rnn(cell_fw, inputs, initial_state=self._initial_state_fw)
      #endoutput = outputs[-1]

    if FLAGS.minibatch_d:
      outputs = [minibatch(tf.reshape(outp, shape=[FLAGS.batch_size, -1]), msg=msg, reuse_scope=(i!=0)) for i,outp in enumerate(outputs)]
    # decision = tf.sigmoid(linear(outputs[-1], 1, 'decision'))
    if FLAGS.end_classification:
      decisions = [tf.sigmoid(linear(output, 1, 'decision', reuse_scope=(i!=0))) for i,output in enumerate([outputs[0], outputs[-1]])]
      decisions = tf.pack(decisions)
      decisions = tf.transpose(decisions, perm=[1,0,2])
      print('shape, decisions: {}'.format(decisions.get_shape()))
    else:
      decisions = [tf.sigmoid(linear(output, 1, 'decision', reuse_scope=(i!=0))) for i,output in enumerate(outputs)]
      decisions = tf.pack(decisions)
      decisions = tf.transpose(decisions, perm=[1,0,2])
      print('shape, decisions: {}'.format(decisions.get_shape()))
    decision = tf.reduce_mean(decisions, reduction_indices=[1,2])
    decision = tf.Print(decision, [decision],
            '{} decision = '.format(msg), summarize=20, first_n=20)
    return (decision,tf.transpose(tf.pack(outputs), perm=[1,0,2]))
      

  
  def assign_lr(self, session, lr_value):
    session.run(self._lr_update, feed_dict={self._new_lr: lr_value})

  @property
  def generated_features(self):
    return self._generated_features

  @property
  def input_songdata(self):
    return self._input_songdata

  @property
  def input_metadata(self):
    return self._input_metadata

  @property
  def targets(self):
    return self._targets

  @property
  def initial_state(self):
    return self._initial_state

  @property
  def cost(self):
    return self._cost

  @property
  def final_state(self):
    return self._final_state

  @property
  def lr(self):
    return self._lr

  @property
  def train_op(self):
    return self._train_op



def run_epoch(session, model, loader, datasetlabel, eval_op_g, eval_op_d, pretraining=False, verbose=False, run_metadata=None, pretraining_d=False):
  """Runs the model on the given data."""
  #epoch_size = ((len(data) // model.batch_size) - 1) // model.songlength
  epoch_start_time = time.time()
  g_loss, d_loss = 10.0, 10.0
  g_losses, d_losses = 0.0, 0.0
  iters = 0
  #state = session.run(model.initial_state)
  time_before_graph = None
  time_after_graph = None
  times_in_graph = []
  times_in_python = []
  #times_in_batchreading = []
  loader.rewind(part=datasetlabel)
  [batch_meta, batch_song] = loader.get_batch(model.batch_size, model.songlength, part=datasetlabel)

  run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)

  while batch_meta is not None and batch_song is not None:
    op_g = eval_op_g
    op_d = eval_op_d
    if datasetlabel == 'train' and not pretraining: # and not FLAGS.feature_matching:
      if d_loss == 0.0 and g_loss == 0.0:
        print('Both G and D train loss are zero. Exiting.')
        break
        #saver.save(session, checkpoint_path, global_step=m.global_step)
        #break
      elif d_loss == 0.0:
        #print('D train loss is zero. Freezing optimization. G loss: {:.3f}'.format(g_loss))
        op_g = tf.no_op()
      elif g_loss == 0.0: 
        #print('G train loss is zero. Freezing optimization. D loss: {:.3f}'.format(d_loss))
        op_d = tf.no_op()
      elif g_loss < 2.0 or d_loss < 2.0:
        if g_loss*.7 > d_loss:
          #print('G train loss is {:.3f}, D train loss is {:.3f}. Freezing optimization of D'.format(g_loss, d_loss))
          op_g = tf.no_op()
        #elif d_loss*.7 > g_loss:
          #print('G train loss is {:.3f}, D train loss is {:.3f}. Freezing optimization of G'.format(g_loss, d_loss))
        op_d = tf.no_op()
    #fetches = [model.cost, model.final_state, eval_op]
    if pretraining:
      if pretraining_d:
        fetches = [model.rnn_pretraining_loss, model.d_loss, op_g, op_d]
      else:
        fetches = [model.rnn_pretraining_loss, tf.no_op(), op_g, op_d]
    else:
      fetches = [model.g_loss, model.d_loss, op_g, op_d]
    feed_dict = {}
    feed_dict[model.input_songdata.name] = batch_song
    feed_dict[model.input_metadata.name] = batch_meta
    #print(batch_song)
    #print(batch_song.shape)
    
    #for i, (c, h) in enumerate(model.initial_state):
    #  feed_dict[c] = state[i].c
    #  feed_dict[h] = state[i].h
    #cost, state, _ = session.run(fetches, feed_dict)
    time_before_graph = time.time()
    if iters > 0:
      times_in_python.append(time_before_graph-time_after_graph)
    if run_metadata:
      g_loss, d_loss, _, _ = session.run(fetches, feed_dict, options=run_options, run_metadata=run_metadata)
    else:
      g_loss, d_loss, _, _ = session.run(fetches, feed_dict)
    time_after_graph = time.time()
    if iters > 0:
      times_in_graph.append(time_after_graph-time_before_graph)
    g_losses += g_loss
    if not pretraining:
      d_losses += d_loss
    iters += 1

    if verbose and iters % 10 == 9:
      songs_per_sec = float(iters * model.batch_size)/float(time.time() - epoch_start_time)
      avg_time_in_graph = float(sum(times_in_graph))/float(len(times_in_graph))
      avg_time_in_python = float(sum(times_in_python))/float(len(times_in_python))
      #avg_time_batchreading = float(sum(times_in_batchreading))/float(len(times_in_batchreading))
      if pretraining:
        print("{}: {} (pretraining) batch loss: G: {:.3f}, avg loss: G: {:.3f}, speed: {:.1f} songs/s, avg in graph: {:.1f}, avg in python: {:.1f}.".format(datasetlabel, iters, g_loss, float(g_losses)/float(iters), songs_per_sec, avg_time_in_graph, avg_time_in_python))
      else:
        print("{}: {} batch loss: G: {:.3f}, D: {:.3f}, avg loss: G: {:.3f}, D: {:.3f} speed: {:.1f} songs/s, avg in graph: {:.1f}, avg in python: {:.1f}.".format(datasetlabel, iters, g_loss, d_loss, float(g_losses)/float(iters), float(d_losses)/float(iters),songs_per_sec, avg_time_in_graph, avg_time_in_python))
    #batchtime = time.time()
    [batch_meta, batch_song] = loader.get_batch(model.batch_size, model.songlength, part=datasetlabel)
    #times_in_batchreading.append(time.time()-batchtime)

  if iters == 0:
    return (None,None)

  g_mean_loss = g_losses/iters
  if pretraining and not pretraining_d:
    d_mean_loss = None
  else:
    d_mean_loss = d_losses/iters
  return (g_mean_loss, d_mean_loss)


def sample(session, model, batch=False):
  """Samples from the generative model."""
  #state = session.run(model.initial_state)
  fetches = [model.generated_features]
  feed_dict = {}
  generated_features, = session.run(fetches, feed_dict)
  #print( generated_features)
  print( generated_features[0].shape)
  # The following worked when batch_size=1.
  # generated_features = [np.squeeze(x, axis=0) for x in generated_features]
  # If batch_size != 1, we just pick the first sample. Wastefull, yes.
  returnable = []
  if batch:
    for batchno in xrange(generated_features[0].shape[0]):
      returnable.append([x[batchno,:] for x in generated_features])
  else:
    returnable = [x[0,:] for x in generated_features]
  return returnable

def main(_):
  if not FLAGS.datadir:
    raise ValueError("Must set --datadir to midi music dir.")
  if not FLAGS.traindir:
    raise ValueError("Must set --traindir to dir where I can save model and plots.")
 
  restore_flags()
 
  summaries_dir = None
  plots_dir = None
  generated_data_dir = None
  summaries_dir = os.path.join(FLAGS.traindir, 'summaries')
  plots_dir = os.path.join(FLAGS.traindir, 'plots')
  generated_data_dir = os.path.join(FLAGS.traindir, 'generated_data')
  try: os.makedirs(FLAGS.traindir)
  except: pass
  try: os.makedirs(summaries_dir)
  except: pass
  try: os.makedirs(plots_dir)
  except: pass
  try: os.makedirs(generated_data_dir)
  except: pass
  directorynames = FLAGS.traindir.split('/')
  experiment_label = ''
  while not experiment_label:
    experiment_label = directorynames.pop()
  
  global_step = -1
  if os.path.exists(os.path.join(FLAGS.traindir, 'global_step.pkl')):
    with open(os.path.join(FLAGS.traindir, 'global_step.pkl'), 'r') as f:
      global_step = pkl.load(f)
  global_step += 1

  songfeatures_filename = os.path.join(FLAGS.traindir, 'num_song_features.pkl')
  metafeatures_filename = os.path.join(FLAGS.traindir, 'num_meta_features.pkl')
  synthetic=None
  if FLAGS.synthetic_chords:
    synthetic = 'chords'
    print('Training on synthetic chords!')
  if FLAGS.composer is not None:
    print('Single composer: {}'.format(FLAGS.composer))
  loader = music_data_utils.MusicDataLoader(FLAGS.datadir, FLAGS.select_validation_percentage, FLAGS.select_test_percentage, FLAGS.works_per_composer, FLAGS.pace_events, synthetic=synthetic, tones_per_cell=FLAGS.tones_per_cell, single_composer=FLAGS.composer)
  if FLAGS.synthetic_chords:
    # This is just a print out, to check the generated data.
    batch = loader.get_batch(batchsize=1, songlength=400)
    loader.get_midi_pattern([batch[1][0][i] for i in xrange(batch[1].shape[1])])

  num_song_features = loader.get_num_song_features()
  print('num_song_features:{}'.format(num_song_features))
  num_meta_features = loader.get_num_meta_features()
  print('num_meta_features:{}'.format(num_meta_features))

  train_start_time = time.time()
  checkpoint_path = os.path.join(FLAGS.traindir, "model.ckpt")

  songlength_ceiling = FLAGS.songlength

  if global_step < FLAGS.pretraining_epochs:
    #FLAGS.songlength = int(min(((global_step+10)/10)*10,songlength_ceiling))
    FLAGS.songlength = int(min((global_step+1)*4,songlength_ceiling))
 
  with tf.Graph().as_default(), tf.Session(config=tf.ConfigProto(log_device_placement=FLAGS.log_device_placement)) as session:
    with tf.variable_scope("model", reuse=None) as scope:
      scope.set_regularizer(tf.contrib.layers.l2_regularizer(scale=FLAGS.reg_scale))
      m = RNNGAN(is_training=True, num_song_features=num_song_features, num_meta_features=num_meta_features)


    if FLAGS.initialize_d:
      vars_to_restore = {}
      for v in tf.trainable_variables():
        if v.name.startswith('model/G/'):
          print(v.name[:-2])
          vars_to_restore[v.name[:-2]] = v
      saver = tf.train.Saver(vars_to_restore)
      ckpt = tf.train.get_checkpoint_state(FLAGS.traindir)
      if ckpt and tf.gfile.Exists(ckpt.model_checkpoint_path):
        print("Reading model parameters from %s" % ckpt.model_checkpoint_path)
        saver.restore(session, ckpt.model_checkpoint_path)
        session.run(tf.initialize_variables([v for v in tf.trainable_variables() if v.name.startswith('model/D/')]))
      else:
        print("Created model with fresh parameters.")
        session.run(tf.initialize_all_variables())
      saver = tf.train.Saver(tf.all_variables())
    else:
      saver = tf.train.Saver(tf.all_variables())
      ckpt = tf.train.get_checkpoint_state(FLAGS.traindir)
      if ckpt and tf.gfile.Exists(ckpt.model_checkpoint_path):
        print("Reading model parameters from %s" % ckpt.model_checkpoint_path)
        saver.restore(session, ckpt.model_checkpoint_path)
      else:
        print("Created model with fresh parameters.")
        session.run(tf.initialize_all_variables())

    run_metadata = None
    if FLAGS.profiling:
      run_metadata = tf.RunMetadata()
    if not FLAGS.sample:
      train_g_loss,train_d_loss = 1.0,1.0
      for i in range(global_step, FLAGS.max_epoch):
        lr_decay = FLAGS.lr_decay ** max(i - FLAGS.epochs_before_decay, 0.0)

        if global_step < FLAGS.pretraining_epochs:
          #new_songlength = int(min(((i+10)/10)*10,songlength_ceiling))
          new_songlength = int(min((i+1)*4,songlength_ceiling))
        else:
          new_songlength = songlength_ceiling
        if new_songlength != FLAGS.songlength:
          print('Changing songlength, now training on {} events from songs.'.format(new_songlength))
          FLAGS.songlength = new_songlength
          with tf.variable_scope("model", reuse=True) as scope:
            scope.set_regularizer(tf.contrib.layers.l2_regularizer(scale=FLAGS.reg_scale))
            m = RNNGAN(is_training=True, num_song_features=num_song_features, num_meta_features=num_meta_features)

        if not FLAGS.adam:
          m.assign_lr(session, FLAGS.learning_rate * lr_decay)

        save = False
        do_exit = False

        print("Epoch: {} Learning rate: {:.3f}, pretraining: {}".format(i, session.run(m.lr), (i<FLAGS.pretraining_epochs)))
        if i<FLAGS.pretraining_epochs:
          opt_d = tf.no_op()
          if FLAGS.pretraining_d:
            opt_d = m.opt_d
          train_g_loss,train_d_loss = run_epoch(session, m, loader, 'train', m.opt_pretraining, opt_d, pretraining = True, verbose=True, run_metadata=run_metadata, pretraining_d=FLAGS.pretraining_d)
          if FLAGS.pretraining_d:
            try:
              print("Epoch: {} Pretraining loss: G: {:.3f}, D: {:.3f}".format(i, train_g_loss, train_d_loss))
            except:
              print(train_g_loss)
              print(train_d_loss)
          else:
            print("Epoch: {} Pretraining loss: G: {:.3f}".format(i, train_g_loss))
        else:
          train_g_loss,train_d_loss = run_epoch(session, m, loader, 'train', m.opt_d, m.opt_g, verbose=True, run_metadata=run_metadata)
          try:
            print("Epoch: {} Train loss: G: {:.3f}, D: {:.3f}".format(i, train_g_loss, train_d_loss))
          except:
            print("Epoch: {} Train loss: G: {}, D: {}".format(i, train_g_loss, train_d_loss))
        valid_g_loss,valid_d_loss = run_epoch(session, m, loader, 'validation', tf.no_op(), tf.no_op())
        try:
          print("Epoch: {} Valid loss: G: {:.3f}, D: {:.3f}".format(i, valid_g_loss, valid_d_loss))
        except:
          print("Epoch: {} Valid loss: G: {}, D: {}".format(i, valid_g_loss, valid_d_loss))
        
        if train_d_loss == 0.0 and train_g_loss == 0.0:
          print('Both G and D train loss are zero. Exiting.')
          save = True
          do_exit = True
        if i % FLAGS.epochs_per_checkpoint == 0:
          save = True
        if FLAGS.exit_after > 0 and time.time() - train_start_time > FLAGS.exit_after*60:
          print("%s: Has been running for %d seconds. Will exit (exiting after %d minutes)."%(datetime.datetime.today().strftime('%Y-%m-%d %H:%M:%S'), (int)(time.time() - train_start_time), FLAGS.exit_after))
          save = True
          do_exit = True

        if save:
          saver.save(session, checkpoint_path, global_step=i)
          with open(os.path.join(FLAGS.traindir, 'global_step.pkl'), 'w') as f:
            pkl.dump(i, f)
          if FLAGS.profiling:
            # Create the Timeline object, and write it to a json
            tl = timeline.Timeline(run_metadata.step_stats)
            ctf = tl.generate_chrome_trace_format()
            with open(os.path.join(plots_dir, 'timeline.json'), 'w') as f:
              f.write(ctf)
          print('{}: Saving done!'.format(i))

        step_time, loss = 0.0, 0.0
        if train_d_loss is None: #pretraining
          train_d_loss = 0.0
          valid_d_loss = 0.0
          valid_g_loss = 0.0
        if not os.path.exists(os.path.join(plots_dir, 'gnuplot-input.txt')):
          with open(os.path.join(plots_dir, 'gnuplot-input.txt'), 'w') as f:
            f.write('# global-step learning-rate train-g-loss train-d-loss valid-g-loss valid-d-loss\n')
        with open(os.path.join(plots_dir, 'gnuplot-input.txt'), 'a') as f:
          try:
            f.write('{} {:.4f} {:.2f} {:.2f} {:.3} {:.3f}\n'.format(i, m.lr.eval(), train_g_loss, train_d_loss, valid_g_loss, valid_d_loss))
          except:
            f.write('{} {} {} {} {} {}\n'.format(i, m.lr.eval(), train_g_loss, train_d_loss, valid_g_loss, valid_d_loss))
        if not os.path.exists(os.path.join(plots_dir, 'gnuplot-commands-loss.txt')):
          with open(os.path.join(plots_dir, 'gnuplot-commands-loss.txt'), 'a') as f:
            f.write('set terminal postscript eps color butt "Times" 14\nset yrange [0:400]\nset output "loss.eps"\nplot \'gnuplot-input.txt\' using ($1):($3) title \'train G\' with linespoints, \'gnuplot-input.txt\' using ($1):($4) title \'train D\' with linespoints, \'gnuplot-input.txt\' using ($1):($5) title \'valid G\' with linespoints, \'gnuplot-input.txt\' using ($1):($6) title \'valid D\' with linespoints, \n')
        if not os.path.exists(os.path.join(plots_dir, 'gnuplot-commands-midistats.txt')):
          with open(os.path.join(plots_dir, 'gnuplot-commands-midistats.txt'), 'a') as f:
            f.write('set terminal postscript eps color butt "Times" 14\nset yrange [0:127]\nset xrange [0:70]\nset output "midistats.eps"\nplot \'midi_stats.gnuplot\' using ($1):(100*$3) title \'Scale consistency, %\' with linespoints, \'midi_stats.gnuplot\' using ($1):($6) title \'Tone span, halftones\' with linespoints, \'midi_stats.gnuplot\' using ($1):($10) title \'Unique tones\' with linespoints, \'midi_stats.gnuplot\' using ($1):($23) title \'Intensity span, units\' with linespoints, \'midi_stats.gnuplot\' using ($1):(100*$24) title \'Polyphony, %\' with linespoints, \'midi_stats.gnuplot\' using ($1):($12) title \'3-tone repetitions\' with linespoints\n')
        try:
          Popen(['gnuplot','gnuplot-commands-loss.txt'], cwd=plots_dir)
          Popen(['gnuplot','gnuplot-commands-midistats.txt'], cwd=plots_dir)
        except:
          print('failed to run gnuplot. Please do so yourself: gnuplot gnuplot-commands.txt cwd={}'.format(plots_dir))
        
        song_data = sample(session, m, batch=True)
        midi_patterns = []
        print('formatting midi...')
        midi_time = time.time()
        for d in song_data:
          midi_patterns.append(loader.get_midi_pattern(d))
        print('done. time: {}'.format(time.time()-midi_time))
        
        filename = os.path.join(generated_data_dir, 'out-{}-{}-{}.mid'.format(experiment_label, i, datetime.datetime.today().strftime('%Y-%m-%d-%H-%M-%S')))
        loader.save_midi_pattern(filename, midi_patterns[0])
  
        stats = []
        print('getting stats...')
        stats_time = time.time()
        for p in midi_patterns:
          stats.append(get_all_stats(p))
        print('done. time: {}'.format(time.time()-stats_time))
        #print(stats)
        stats = [stat for stat in stats if stat is not None]
        if len(stats):
          stats_keys_string = ['scale']
          stats_keys = ['scale_score', 'tone_min', 'tone_max', 'tone_span', 'freq_min', 'freq_max', 'freq_span', 'tones_unique', 'repetitions_2', 'repetitions_3', 'repetitions_4', 'repetitions_5', 'repetitions_6', 'repetitions_7', 'repetitions_8', 'repetitions_9', 'estimated_beat', 'estimated_beat_avg_ticks_off', 'intensity_min', 'intensity_max', 'intensity_span', 'polyphony_score', 'top_2_interval_difference', 'top_3_interval_difference', 'num_tones']
          statsfilename = os.path.join(plots_dir, 'midi_stats.gnuplot')
          if not os.path.exists(statsfilename):
            with open(statsfilename, 'a') as f:
              f.write('# Average numers over one minibatch of size {}.\n'.format(FLAGS.batch_size))
              f.write('# global-step {} {}\n'.format(' '.join([s.replace(' ', '_') for s in stats_keys_string]), ' '.join(stats_keys)))
          with open(statsfilename, 'a') as f:
            f.write('{} {} {}\n'.format(i, ' '.join(['{}'.format(stats[0][key].replace(' ', '_')) for key in stats_keys_string]), ' '.join(['{:.3f}'.format(sum([s[key] for s in stats])/float(len(stats))) for key in stats_keys])))
          print('Saved {}.'.format(filename))
          
        if do_exit:
          if FLAGS.call_after is not None:
            print("%s: Will call \"%s\" before exiting."%(datetime.datetime.today().strftime('%Y-%m-%d %H:%M:%S'), FLAGS.call_after))
            res = call(FLAGS.call_after.split(" "))
            print ('{}: call returned {}.'.format(datetime.datetime.today().strftime('%Y-%m-%d %H:%M:%S'), res))
          exit()
        sys.stdout.flush()


      test_g_loss,test_d_loss = run_epoch(session, m, loader, 'test', tf.no_op(), tf.no_op())
      print("Test loss G: %.3f, D: %.3f" %(test_g_loss, test_d_loss))

    song_data = sample(session, m)
    filename = os.path.join(generated_data_dir, 'out-{}-{}-{}.mid'.format(experiment_label, i, datetime.datetime.today().strftime('%Y-%m-%d-%H-%M-%S')))
    loader.save_data(filename, song_data)
    print('Saved {}.'.format(filename))



if __name__ == "__main__":
  tf.app.run()