53 lines
1.6 KiB
Python
53 lines
1.6 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
|
|
import numpy as np
|
|
import sys
|
|
import pickle
|
|
|
|
train_data_path_X = sys.argv[1]
|
|
train_data_path_y = sys.argv[2]
|
|
save_model_path = sys.argv[3]
|
|
epochs = int(sys.argv[4])
|
|
|
|
# model architecture
|
|
# model = Sequential()
|
|
# model.add(LSTM(128, activation='relu', input_shape=(96,128)))
|
|
# model.add(RepeatVector(96))
|
|
# model.add(LSTM(128, activation='softmax', return_sequences=True))
|
|
# model.add(TimeDistributed(Dense(128)))
|
|
#
|
|
# model.compile(optimizer='adam',
|
|
# loss='categorical_crossentropy',
|
|
# metrics=['accuracy'])
|
|
|
|
model = Sequential()
|
|
model.add(LSTM(128,input_shape=(96, 128),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(512))
|
|
# model.add(Dropout(0.3))
|
|
# model.add(Dense(128))
|
|
model.add(Activation('softmax'))
|
|
model.add(Reshape((4, 128)))
|
|
model.compile(loss='categorical_crossentropy', optimizer='rmsprop')
|
|
|
|
# load training data
|
|
print('Reading samples from: {}'.format(train_data_path_X))
|
|
train_X = np.load(train_data_path_X)['arr_0']
|
|
train_y = np.load(train_data_path_y)['arr_0']
|
|
|
|
# model training
|
|
model.fit(train_X, train_y, epochs=epochs, batch_size=32)
|
|
|
|
# save trained model
|
|
pickle_path = '{}.pickle'.format(save_model_path)
|
|
pickle.dump(model, open(pickle_path,'wb'))
|
|
print("Model save to {}".format(pickle_path))
|