ENSEMBLE_SIZE = 5
parser = ArgumentParser()
parser.add_argument('--mode', default='train')
parser.add_argument('--pause', default=0)
args = parser.parse_args()
task = HateSpeech()
vocab_size = task.max_vocab_indexes['syllable_contents']
models = []
for i in range(ENSEMBLE_SIZE):
model = BaseLine(HIDDEN_DIM, DROPOUT_RATE, vocab_size, EMBEDDING_SIZE, BI_RNN_LAYERS, UNI_RNN_LAYERS)
model.to('cuda')
models.append(model)
if args.pause:
bind_models(models)
nsml.paused(scope=locals())
if args.mode == 'train':
scores = []
for i, model in enumerate(models):
trainer = Trainer(model, i, ENSEMBLE_SIZE, device='cuda')
best_state_dict, best_f1_score = trainer.train()
model.load_state_dict(best_state_dict)
scores.append(best_f1_score)
print('best f1 score:', best_f1_score)
print('avg f1 score:', sum(scores) / len(scores))
bind_models(models)
nsml.save('ensemble')
