praca-magisterska/project/train.py

69 lines
2.5 KiB
Python

#!python3
#!/usr/bin/env python3
import sys
import pickle
import settings
import numpy as np
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
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_X = np.unique(train_X)
vocab_y = np.unique(train_y)
vocab = np.concatenate([vocab_X, vocab_y])
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')
# This code will train our model, with given by parameter number of epochs
print('Training...')
model.fit(train_X, train_y, epochs=epochs, batch_size=64)
# it saves model, and additional informations of model
# that is needed to generate music from it
pickle.dump(model, open(save_model_path,'wb'))
pickle.dump((int_to_note, n_vocab, train_X.shape[1]), open('{}_dict'.format(save_model_path),'wb'))
print('Done!')
print("Model saved to: {}".format(save_model_path))