Refactored into train and generate.py

Author: calclavia

dataset.py

@@ -9,7 +9,7 @@
 
 from constants import *
 from midi_util import load_midi
-from util import chunk, get_all_files, one_hot
+from util import *
 
 def compute_beat(beat, notes_in_bar):
     return one_hot(beat % notes_in_bar, notes_in_bar)

generate.py

@@ -1,8 +1,11 @@
 import numpy as np
+import tensorflow as tf
 from collections import deque
 import midi
+import argparse
 
 from constants import *
+from util import *
 from dataset import *
 from tqdm import tqdm
 from midi_util import midi_encode
@@ -90,7 +93,7 @@ def process_inputs(ins):
     ins = [np.array(i) for i in ins]
     return ins
 
-def generate(models, num_bars=16, styles=[0, 4, 12, 20]):
+def generate(models, num_bars=16, styles=[[1,0,0,0]]):
     print('Generating with styles:', styles)
 
     _, time_model, note_model = models
@@ -115,14 +118,26 @@ def generate(models, num_bars=16, styles=[0, 4, 12, 20]):
         # Move one time step
         yield [g.end_time(t) for g in generations]
 
-def write_file(name, results):
+def write_file(fpath, results):
     """
     Takes a list of all notes generated per track and writes it to file
     """
     results = zip(*list(results))
 
     for i, result in enumerate(results):
-        fpath = SAMPLES_DIR + '/' + name + '_' + str(i) + '.mid'
         print('Writing file', fpath)
+        os.makedirs(os.path.dirname(fpath), exist_ok=True)
         mf = midi_encode(unclamp_midi(result))
         midi.write_midifile(fpath, mf)
+
+def main():
+    # parser = argparse.ArgumentParser(description='Generates music.')
+    # parser.add_argument('--gen', default=None, nargs='+', help='Style to generate')
+    # args = parser.parse_args()
+
+    with tf.device('cpu:0'):
+        models = build_or_load()
+        write_file(os.path.join(SAMPLES_DIR, 'output.mid'), generate(models))
+
+if __name__ == '__main__':
+    main()

model.py

@@ -46,6 +46,46 @@ def f(x):
         return bins
     return f
 
+def time_axis(dropout):
+    def f(notes, beat, style):
+        time_steps = int(notes.get_shape()[1])
+        
+        # TODO: Experiment with when to apply conv
+        note_octave = TimeDistributed(Conv1D(OCTAVE_UNITS, 2 * OCTAVE, padding='same'))(notes)
+        note_octave = Activation('tanh')(note_octave)
+        note_octave = Dropout(dropout)(note_octave)
+
+        # Create features for every single note.
+        note_features = Concatenate()([
+            Lambda(pitch_pos_in_f(time_steps))(notes),
+            Lambda(pitch_class_in_f(time_steps))(notes),
+            Lambda(pitch_bins_f(time_steps))(notes),
+            note_octave,
+            TimeDistributed(RepeatVector(NUM_NOTES))(beat)
+        ])
+
+        x = note_features
+
+        # [batch, notes, time, features]
+        x = Permute((2, 1, 3))(x)
+
+        # Apply LSTMs
+        for l in range(TIME_AXIS_LAYERS):
+            # Integrate style
+            style_proj = Dense(int(x.get_shape()[3]))(style)
+            style_proj = Activation('tanh')(style_proj)
+            style_proj = Dropout(dropout)(style_proj)
+            style_proj = TimeDistributed(RepeatVector(NUM_NOTES))(style_proj)
+            style_proj = Permute((2, 1, 3))(style_proj)
+            x = Add()([x, style_proj])
+
+            x = TimeDistributed(LSTM(TIME_AXIS_UNITS, return_sequences=True))(x)
+            x = Dropout(dropout)(x)
+
+        # [batch, time, notes, features]
+        return Permute((2, 1, 3))(x)
+    return f
+
 def note_axis(dropout):
     def f(x, chosen, style):
         time_steps = int(x.get_shape()[1])
@@ -74,7 +114,7 @@ def f(x, chosen, style):
     return f
 
 def style_layer(input_dropout):
-    emb = Embedding(NUM_STYLES, STYLE_UNITS)
+    emb = Dense(STYLE_UNITS)
     def f(style_in):
         style = emb(style_in)
         return Dropout(input_dropout)(style)
@@ -83,7 +123,7 @@ def f(style_in):
 def build_models(time_steps=SEQ_LEN, input_dropout=0.2, dropout=0.5):
     notes_in = Input((time_steps, NUM_NOTES, NOTE_UNITS))
     beat_in = Input((time_steps, NOTES_PER_BAR))
-    style_in = Input((time_steps,))
+    style_in = Input((time_steps, NUM_STYLES))
     # Target input for conditioning
     chosen_in = Input((time_steps, NUM_NOTES, NOTE_UNITS))
 
@@ -97,40 +137,7 @@ def build_models(time_steps=SEQ_LEN, input_dropout=0.2, dropout=0.5):
     style = style_l(style_in)
 
     """ Time axis """
-    # TODO: Experiment with when to apply conv
-    note_octave = TimeDistributed(Conv1D(OCTAVE_UNITS, 2 * OCTAVE, padding='same'))(notes)
-    note_octave = Activation('tanh')(note_octave)
-    note_octave = Dropout(dropout)(note_octave)
-
-    # Create features for every single note.
-    note_features = Concatenate()([
-        Lambda(pitch_pos_in_f(time_steps))(notes),
-        Lambda(pitch_class_in_f(time_steps))(notes),
-        Lambda(pitch_bins_f(time_steps))(notes),
-        note_octave,
-        TimeDistributed(RepeatVector(NUM_NOTES))(beat)
-    ])
-
-    x = note_features
-
-    # [batch, notes, time, features]
-    x = Permute((2, 1, 3))(x)
-
-    # Apply LSTMs
-    for l in range(TIME_AXIS_LAYERS):
-        # Integrate style
-        style_proj = Dense(int(x.get_shape()[3]))(style)
-        style_proj = Activation('tanh')(style_proj)
-        style_proj = Dropout(dropout)(style_proj)
-        style_proj = TimeDistributed(RepeatVector(NUM_NOTES))(style_proj)
-        style_proj = Permute((2, 1, 3))(style_proj)
-        x = Add()([x, style_proj])
-
-        x = TimeDistributed(LSTM(TIME_AXIS_UNITS, return_sequences=True))(x)
-        x = Dropout(dropout)(x)
-
-    # [batch, time, notes, features]
-    time_out = Permute((2, 1, 3))(x)
+    time_out = time_axis(dropout)(notes, beat, style)
 
     """ Note Axis & Prediction Layer """
     naxis = note_axis(dropout)
@@ -144,7 +151,7 @@ def build_models(time_steps=SEQ_LEN, input_dropout=0.2, dropout=0.5):
 
     note_features = Input((1, NUM_NOTES, TIME_AXIS_UNITS), name='note_features')
     chosen_gen_in = Input((1, NUM_NOTES, NOTE_UNITS), name='chosen_gen_in')
-    style_gen_in = Input((1,), name='style_in')
+    style_gen_in = Input((1, NUM_STYLES), name='style_in')
 
     # Dropout inputs
     chosen_gen = Dropout(input_dropout)(chosen_gen_in)

main.py renamed to train.py

@@ -13,27 +13,11 @@
 
 def main():
     parser = argparse.ArgumentParser(description='Generates music.')
-    parser.add_argument('--train', default=False, action='store_true', help='Train model?')
     parser.add_argument('--gen', default=False, action='store_true', help='Generate after each epoch?')
     args = parser.parse_args()
 
     models = build_or_load()
-
-    if args.train:
-        train(models, args.gen)
-    else:
-        write_file(os.path.join(SAMPLES_DIR, 'output.mid'), generate(models))
-
-def build_or_load(allow_load=True):
-    models = build_models()
-    models[0].summary()
-    if allow_load:
-        try:
-            models[0].load_weights(MODEL_FILE)
-            print('Loaded model from file.')
-        except:
-            print('Unable to load model from file.')
-    return models
+    train(models, args.gen)
 
 def train(models, gen):
     print('Loading data')

util.py

@@ -10,15 +10,17 @@ def one_hot(i, nb_classes):
     arr[i] = 1
     return arr
 
-def chunk(a, size):
-    # Zero pad extra spaces
-    target_size = math.ceil(len(a) / float(size)) * size
-    inc_size = target_size - len(a)
-    assert inc_size >= 0 and inc_size < size, inc_size
-    a = np.array(a)
-    a = np.pad(a, [(0, inc_size)] + [(0, 0) for i in range(len(a.shape) - 1)], mode='constant')
-    assert a.shape[0] == target_size
-    return np.swapaxes(np.split(a, size), 0, 1)
+def build_or_load(allow_load=True):
+    from model import build_models
+    models = build_models()
+    models[0].summary()
+    if allow_load:
+        try:
+            models[0].load_weights(MODEL_FILE)
+            print('Loaded model from file.')
+        except:
+            print('Unable to load model from file.')
+    return models
 
 def get_all_files(paths):
     potential_files = []