Flag to control whether or not loaded/pretrained model is re-trained on eval task
@@ -48,6 +48,8 @@ def main(arguments):
type=str)
parser.add_argument('--eval_tasks', help='list of additional tasks to train a classifier,' +
'then evaluate on', type=str, default='')
+ parser.add_argument('--train_for_eval', help='1 if models should be trained for the eval tasks (defaults to True)',
+ type=int, default=1)
parser.add_argument('--classifier', help='type of classifier to use', type=str,
default='log_reg', choices=['log_reg', 'mlp', 'fancy_mlp'])
parser.add_argument('--classifier_hid_dim', help='hid dim of classifier', type=int, default=512)
@@ -216,6 +218,7 @@ def main(arguments):
# Train just the task-specific components for eval tasks
# TODO(Alex): currently will overwrite model checkpoints from training
for task in eval_tasks:
+ if args.train_for_eval:
pred_module = getattr(model, "%s_mdl" % task.name)
to_train = [(n, p) for n, p in pred_module.named_parameters() if p.requires_grad]
trainer, _, opt_params, schd_params = build_trainer(args, model,
