mirror of https://github.com/coqui-ai/TTS.git
125 lines
3.9 KiB
Python
125 lines
3.9 KiB
Python
import argparse
|
|
import pickle
|
|
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
from clearml import Task
|
|
from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
|
|
from sklearn.model_selection import KFold, cross_val_score
|
|
|
|
|
|
def load_data(train, labels):
|
|
t = np.load(train)
|
|
l = np.load(labels)
|
|
return t,l
|
|
|
|
def compute_metrics(y_pred, y_true):
|
|
acc = accuracy_score(y_true.argmax(axis=1),y_pred.argmax(axis=1))
|
|
print(f"Model Accuracy on untrained data: {acc}")
|
|
conf_matrix = confusion_matrix(y_true.argmax(axis=1),y_pred.argmax(axis=1))
|
|
print(f"Confusion matrix on {y_true.shape[-1]} classes: \n {conf_matrix}")
|
|
f1 = f1_score(y_true.argmax(axis=1),y_pred.argmax(axis=1),average='weighted')
|
|
print(f"Weighted f1 score: {f1}")
|
|
|
|
def test_model(test,test_labels,model):
|
|
test, y_true = load_data(test, test_labels)
|
|
y_pred = model.predict(test, verbose=1)
|
|
compute_metrics(y_pred,y_true)
|
|
|
|
def baseline_model(class_num):
|
|
model = tf.keras.models.Sequential()
|
|
model.add(tf.keras.layers.Dense(1024, input_dim=512, activation='relu'))
|
|
model.add(tf.keras.layers.Dropout(0.2))
|
|
model.add(tf.keras.layers.Dense(512, activation='relu'))
|
|
model.add(tf.keras.layers.Dropout(0.5))
|
|
model.add(tf.keras.layers.Dense(128, activation='relu'))
|
|
model.add(tf.keras.layers.Dense(64, activation='relu'))
|
|
model.add(tf.keras.layers.Dense(class_num, activation='softmax'))
|
|
opt = tf.keras.optimizers.Adam(learning_rate=0.0001)
|
|
model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])
|
|
return model
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"train_data",
|
|
nargs='?',
|
|
)
|
|
parser.add_argument(
|
|
"train_labels",
|
|
nargs='?',
|
|
)
|
|
parser.add_argument(
|
|
"test_data",
|
|
nargs='?',
|
|
)
|
|
parser.add_argument(
|
|
"test_labels",
|
|
nargs='?',
|
|
)
|
|
parser.add_argument(
|
|
'--eval',
|
|
action='store_true'
|
|
)
|
|
parser.add_argument(
|
|
'--load_ckpt',
|
|
required=False,
|
|
)
|
|
parser.add_argument(
|
|
'--predict',
|
|
action='store_true'
|
|
)
|
|
parser.add_argument(
|
|
'class_weights',
|
|
)
|
|
parser.add_argument(
|
|
'epoch',
|
|
)
|
|
parser.add_argument(
|
|
"--clearml_project",
|
|
default="YourTTS-sprint2",
|
|
)
|
|
parser.add_argument(
|
|
"--clearml_task",
|
|
default="attribute-classifier",
|
|
)
|
|
parser.add_argument(
|
|
"label",
|
|
)
|
|
|
|
|
|
args = parser.parse_args()
|
|
task = Task.init(
|
|
project_name=args.clearml_project,
|
|
task_name=f"{args.clearml_task}-{args.label}",
|
|
)
|
|
train, labels = load_data(args.train_data, args.train_labels)
|
|
model = baseline_model(len(np.unique(labels,axis=0)))
|
|
with open(args.class_weights, 'rb') as f:
|
|
class_weights = pickle.load(f)
|
|
if not args.eval and not args.predict:
|
|
print(f"Training on: {train.shape} and testing on {labels.shape}")
|
|
EPOCHS = int(args.epoch)
|
|
checkpoint_filepath = "checkpoints/checkpoint"
|
|
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
|
|
filepath=checkpoint_filepath,
|
|
save_weights_only=True,
|
|
monitor='val_accuracy',
|
|
verbose=1,
|
|
mode='max',
|
|
save_best_only=True)
|
|
model.fit(epochs=EPOCHS,x=train,y=labels,validation_split=0.2,batch_size=100, callbacks=[model_checkpoint_callback], class_weight=class_weights)
|
|
else:
|
|
model.load_weights(args.load_ckpt)
|
|
test_model(args.test_data, args.test_labels,model)
|
|
if(args.predict):
|
|
test = np.load(test)
|
|
y_pred = model.predict(test, verbose=1)
|
|
num_labels = np.where(y_pred==1)[1]
|
|
np.save("predicted_labels.npy",num_labels)
|
|
if __name__ == "__main__":
|
|
main()
|
|
# kfold = KFold(n_splits=10, shuffle=True)
|
|
# results = cross_val_score(estimator, train, labels, cv=kfold)
|
|
# print("Baseline: %.2f%% (%.2f%%)" % (results.mean()*100, results.std()*100))
|