Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
75 lines (57 sloc) 2.75 KB
# from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
import numpy as np
import argparse
import keras
from keras.callbacks import ModelCheckpoint
from model import RNet
from data import BatchGen, load_dataset
import sys
sys.setrecursionlimit(100000)
np.random.seed(10)
parser = argparse.ArgumentParser()
parser.add_argument('--hdim', default=75, help='Model to evaluate', type=int)
parser.add_argument('--batch_size', default=70, help='Batch size', type=int)
parser.add_argument('--nb_epochs', default=50, help='Number of Epochs', type=int)
parser.add_argument('--optimizer', default='Adadelta', help='Optimizer', type=str)
parser.add_argument('--lr', default=None, help='Learning rate', type=float)
parser.add_argument('--name', default='', help='Model dump name prefix', type=str)
parser.add_argument('--loss', default='categorical_crossentropy', help='Loss', type=str)
parser.add_argument('--dropout', default=0, type=float)
parser.add_argument('--char_level_embeddings', action='store_true')
parser.add_argument('--train_data', default='data/train_data.pkl', help='Train Set', type=str)
parser.add_argument('--valid_data', default='data/valid_data.pkl', help='Validation Set', type=str)
# parser.add_argument('model', help='Model to evaluate', type=str)
args = parser.parse_args()
print('Creating the model...', end='')
model = RNet(hdim=args.hdim, dropout_rate=args.dropout, N=None, M=None,
char_level_embeddings=args.char_level_embeddings)
print('Done!')
print('Compiling Keras model...', end='')
optimizer_config = {'class_name': args.optimizer,
'config': {'lr': args.lr} if args.lr else {}}
model.compile(optimizer=optimizer_config,
loss=args.loss,
metrics=['accuracy'])
print('Done!')
print('Loading datasets...', end='')
train_data = load_dataset(args.train_data)
valid_data = load_dataset(args.valid_data)
print('Done!')
print('Preparing generators...', end='')
maxlen = [300, 300, 30, 30] if args.char_level_embeddings else [300, 30]
train_data_gen = BatchGen(*train_data, batch_size=args.batch_size, shuffle=False, group=True, maxlen=maxlen)
valid_data_gen = BatchGen(*valid_data, batch_size=args.batch_size, shuffle=False, group=True, maxlen=maxlen)
print('Done!')
print('Training...', end='')
path = 'models/' + args.name + '{epoch}-t{loss}-v{val_loss}.model'
model.fit_generator(generator=train_data_gen,
steps_per_epoch=train_data_gen.steps(),
validation_data=valid_data_gen,
validation_steps=valid_data_gen.steps(),
epochs=args.nb_epochs,
callbacks=[
ModelCheckpoint(path, verbose=1, save_best_only=True)
])
print('Done!')