67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import tensorflow as tf
|
|
import settings
|
|
from tensorflow.keras import layers
|
|
from keras.layers import Input, Dense, Conv2D, Flatten, LSTM, Dropout, TimeDistributed, RepeatVector, Activation, Bidirectional, Reshape
|
|
from keras.models import Model, Sequential
|
|
from keras.utils.np_utils import to_categorical
|
|
import numpy as np
|
|
import sys
|
|
import pickle
|
|
|
|
def load_data(samples_path):
|
|
print('Loading... {}'.format(train_data_path))
|
|
train_X = np.load(train_data_path, allow_pickle=True)['arr_0']
|
|
train_y = np.load(train_data_path, allow_pickle=True)['arr_1']
|
|
return train_X, train_y
|
|
|
|
# TODO: make transformer class with fit, transform and reverse definitions
|
|
def preprocess_samples(train_X, train_y):
|
|
vocab = np.unique(train_X)
|
|
n_vocab = vocab.shape[0]
|
|
note_to_int = dict((note, number) for number, note in enumerate(vocab))
|
|
int_to_note = dict((number, note) for number, note in enumerate(vocab))
|
|
|
|
_train_X = []
|
|
_train_y = []
|
|
for sample in train_X:
|
|
# TODO: add normalizasion
|
|
_train_X.append([note_to_int[note] for note in sample])
|
|
|
|
train_X = np.array(_train_X).reshape(train_X.shape[0], train_X.shape[1], 1)
|
|
train_y = np.array([note_to_int[note] for note in train_y]).reshape(-1,1)
|
|
train_y = to_categorical(train_y)
|
|
|
|
return train_X, train_y, n_vocab, int_to_note
|
|
|
|
train_data_path = sys.argv[1]
|
|
|
|
train_X, train_y = load_data(train_data_path)
|
|
train_X, train_y, n_vocab, int_to_note = preprocess_samples(train_X, train_y)
|
|
|
|
save_model_path = sys.argv[2]
|
|
epochs = int(sys.argv[3])
|
|
|
|
model = Sequential()
|
|
model.add(LSTM(512, input_shape=(train_X.shape[1], train_X.shape[2]), return_sequences=True))
|
|
model.add(Dropout(0.3))
|
|
model.add(LSTM(512, return_sequences=True))
|
|
model.add(Dropout(0.3))
|
|
model.add(LSTM(512))
|
|
model.add(Dense(256))
|
|
model.add(Dropout(0.3))
|
|
model.add(Dense(n_vocab))
|
|
model.add(Activation('softmax'))
|
|
model.compile(loss='categorical_crossentropy', optimizer='rmsprop')
|
|
|
|
# model training
|
|
print('Training...')
|
|
model.fit(train_X, train_y, epochs=epochs, batch_size=64)
|
|
|
|
# save trained model
|
|
pickle.dump(model, open(save_model_path,'wb'))
|
|
pickle.dump(int_to_note, open('{}_dict'.format(save_model_path),'wb'))
|
|
print('Done!')
|
|
print("Model saved to: {}".format(save_model_path))
|