import tensorflow as tf
import numpy as np
import miditoolkit
import modules
import pickle
import utils
import time
class PopMusicTransformer(object):
########################################
# initialize
########################################
def __init__(self, checkpoint, is_training=False):
# load dictionary
self.dictionary_path = '{}/dictionary.pkl'.format(checkpoint)
self.event2word, self.word2event = pickle.load(open(self.dictionary_path, 'rb'))
# model settings
self.x_len = 512
self.mem_len = 512
self.n_layer = 12
self.d_embed = 512
self.d_model = 512
self.dropout = 0.1
self.n_head = 8
self.d_head = self.d_model // self.n_head
self.d_ff = 2048
self.n_token = len(self.event2word)
self.learning_rate = 0.0002
# load model
self.is_training = is_training
if self.is_training:
self.batch_size = 4
else:
self.batch_size = 1
self.checkpoint_path = '{}/model'.format(checkpoint)
self.load_model()
########################################
# load model
########################################
def load_model(self):
# placeholders
self.x = tf.compat.v1.placeholder(tf.int32, shape=[self.batch_size, None])
self.y = tf.compat.v1.placeholder(tf.int32, shape=[self.batch_size, None])
self.mems_i = [tf.compat.v1.placeholder(tf.float32, [self.mem_len, self.batch_size, self.d_model]) for _ in range(self.n_layer)]
# model
self.global_step = tf.compat.v1.train.get_or_create_global_step()
initializer = tf.compat.v1.initializers.random_normal(stddev=0.02, seed=None)
proj_initializer = tf.compat.v1.initializers.random_normal(stddev=0.01, seed=None)
with tf.compat.v1.variable_scope(tf.compat.v1.get_variable_scope()):
xx = tf.transpose(self.x, [1, 0])
yy = tf.transpose(self.y, [1, 0])
loss, self.logits, self.new_mem = modules.transformer(
dec_inp=xx,
target=yy,
mems=self.mems_i,
n_token=self.n_token,
n_layer=self.n_layer,
d_model=self.d_model,
d_embed=self.d_embed,
n_head=self.n_head,
d_head=self.d_head,
d_inner=self.d_ff,
dropout=self.dropout,
dropatt=self.dropout,
initializer=initializer,
proj_initializer=proj_initializer,
is_training=self.is_training,
mem_len=self.mem_len,
cutoffs=[],
div_val=-1,
tie_projs=[],
same_length=False,
clamp_len=-1,
input_perms=None,
target_perms=None,
head_target=None,
untie_r=False,
proj_same_dim=True)
self.avg_loss = tf.reduce_mean(loss)
# vars
all_vars = tf.compat.v1.trainable_variables()
grads = tf.gradients(self.avg_loss, all_vars)
grads_and_vars = list(zip(grads, all_vars))
all_trainable_vars = tf.reduce_sum([tf.reduce_prod(v.shape) for v in tf.compat.v1.trainable_variables()])
# optimizer
decay_lr = tf.compat.v1.train.cosine_decay(
self.learning_rate,
global_step=self.global_step,
decay_steps=400000,
alpha=0.004)
optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=decay_lr)
self.train_op = optimizer.apply_gradients(grads_and_vars, self.global_step)
# saver
self.saver = tf.compat.v1.train.Saver()
config = tf.compat.v1.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
self.sess = tf.compat.v1.Session(config=config)
self.saver.restore(self.sess, self.checkpoint_path)
########################################
# temperature sampling
########################################
def temperature_sampling(self, logits, temperature, topk):
probs = np.exp(logits / temperature) / np.sum(np.exp(logits / temperature))
if topk == 1:
prediction = np.argmax(probs)
else:
sorted_index = np.argsort(probs)[::-1]
candi_index = sorted_index[:topk]
candi_probs = [probs[i] for i in candi_index]
# normalize probs
candi_probs /= sum(candi_probs)
# choose by predicted probs
prediction = np.random.choice(candi_index, size=1, p=candi_probs)[0]
return prediction
########################################
# extract events for prompt continuation
########################################
def extract_events(self, input_path):
note_items, tempo_items = utils.read_items(input_path)
note_items = utils.quantize_items(note_items)
max_time = note_items[-1].end
if 'chord' in self.checkpoint_path:
chord_items = utils.extract_chords(note_items)
items = chord_items + tempo_items + note_items
else:
items = tempo_items + note_items
groups = utils.group_items(items, max_time)
events = utils.item2event(groups)
return events
########################################
# generate
########################################
def generate(self, n_target_bar, temperature, topk, output_path, prompt=None):
# if prompt, load it. Or, random start
if prompt:
events = self.extract_events(prompt)
words = [[self.event2word['{}_{}'.format(e.name, e.value)] for e in events]]
words[0].append(self.event2word['Bar_None'])
else:
words = []
for _ in range(self.batch_size):
ws = [self.event2word['Bar_None']]
if 'chord' in self.checkpoint_path:
tempo_classes = [v for k, v in self.event2word.items() if 'Tempo Class' in k]
tempo_values = [v for k, v in self.event2word.items() if 'Tempo Value' in k]
chords = [v for k, v in self.event2word.items() if 'Chord' in k]
ws.append(self.event2word['Position_1/16'])
ws.append(np.random.choice(chords))
ws.append(self.event2word['Position_1/16'])
ws.append(np.random.choice(tempo_classes))
ws.append(np.random.choice(tempo_values))
else:
tempo_classes = [v for k, v in self.event2word.items() if 'Tempo Class' in k]
tempo_values = [v for k, v in self.event2word.items() if 'Tempo Value' in k]
ws.append(self.event2word['Position_1/16'])
ws.append(np.random.choice(tempo_classes))
ws.append(np.random.choice(tempo_values))
words.append(ws)
# initialize mem
batch_m = [np.zeros((self.mem_len, self.batch_size, self.d_model), dtype=np.float32) for _ in range(self.n_layer)]
# generate
original_length = len(words[0])
initial_flag = 1
current_generated_bar = 0
while current_generated_bar < n_target_bar:
# input
if initial_flag:
temp_x = np.zeros((self.batch_size, original_length))
for b in range(self.batch_size):
for z, t in enumerate(words[b]):
temp_x[b][z] = t
initial_flag = 0
else:
temp_x = np.zeros((self.batch_size, 1))
for b in range(self.batch_size):
temp_x[b][0] = words[b][-1]
# prepare feed dict
feed_dict = {self.x: temp_x}
for m, m_np in zip(self.mems_i, batch_m):
feed_dict[m] = m_np
# model (prediction)
_logits, _new_mem = self.sess.run([self.logits, self.new_mem], feed_dict=feed_dict)
# sampling
_logit = _logits[-1, 0]
word = self.temperature_sampling(
logits=_logit,
temperature=temperature,
topk=topk)
words[0].append(word)
# if bar event (only work for batch_size=1)
if word == self.event2word['Bar_None']:
current_generated_bar += 1
# re-new mem
batch_m = _new_mem
# write
if prompt:
utils.write_midi(
words=words[0][original_length:],
word2event=self.word2event,
output_path=output_path,
prompt_path=prompt)
else:
utils.write_midi(
words=words[0],
word2event=self.word2event,
output_path=output_path,
prompt_path=None)
########################################
# prepare training data
########################################
def prepare_data(self, midi_paths):
# extract events
all_events = []
for path in midi_paths:
events = self.extract_events(path)
all_events.append(events)
# event to word
all_words = []
for events in all_events:
words = []
for event in events:
e = '{}_{}'.format(event.name, event.value)
if e in self.event2word:
words.append(self.event2word[e])
else:
# OOV
if event.name == 'Note Velocity':
# replace with max velocity based on our training data
words.append(self.event2word['Note Velocity_21'])
else:
# something is wrong
# you should handle it for your own purpose
print('something is wrong! {}'.format(e))
all_words.append(words)
# to training data
self.group_size = 5
segments = []
for words in all_words:
pairs = []
for i in range(0, len(words)-self.x_len-1, self.x_len):
x = words[i:i+self.x_len]
y = words[i+1:i+self.x_len+1]
pairs.append([x, y])
pairs = np.array(pairs)
# abandon the last
for i in np.arange(0, len(pairs)-self.group_size, self.group_size*2):
data = pairs[i:i+self.group_size]
if len(data) == self.group_size:
segments.append(data)
segments = np.array(segments)
return segments
########################################
# finetune
########################################
def finetune(self, training_data, output_checkpoint_folder):
# shuffle
index = np.arange(len(training_data))
np.random.shuffle(index)
training_data = training_data[index]
num_batches = len(training_data) // self.batch_size
st = time.time()
for e in range(200):
total_loss = []
for i in range(num_batches):
segments = training_data[self.batch_size*i:self.batch_size*(i+1)]
batch_m = [np.zeros((self.mem_len, self.batch_size, self.d_model), dtype=np.float32) for _ in range(self.n_layer)]
for j in range(self.group_size):
batch_x = segments[:, j, 0, :]
batch_y = segments[:, j, 1, :]
# prepare feed dict
feed_dict = {self.x: batch_x, self.y: batch_y}
for m, m_np in zip(self.mems_i, batch_m):
feed_dict[m] = m_np
# run
_, gs_, loss_, new_mem_ = self.sess.run([self.train_op, self.global_step, self.avg_loss, self.new_mem], feed_dict=feed_dict)
batch_m = new_mem_
total_loss.append(loss_)
print('>>> Epoch: {}, Step: {}, Loss: {:.5f}, Time: {:.2f}'.format(e, gs_, loss_, time.time()-st))
self.saver.save(self.sess, '{}/model-{:03d}-{:.3f}'.format(output_checkpoint_folder, e, np.mean(total_loss)))
# stop
if np.mean(total_loss) <= 0.1:
break
########################################
# close
########################################
def close(self):
self.sess.close()