def test_mmtl_singletask(self): """One task with one train payload and one labelset""" N = 600 T = 1 tasks = create_tasks(T) model = MetalModel(tasks, verbose=False) payloads = create_payloads(N, T, batch_size=2) metrics_dict = self.trainer.train_model(model, payloads) self.assertEqual(len(metrics_dict), len(SPLITS) * T) for metric, score in metrics_dict.items(): self.assertGreater(score, 0.9)
def test_mmtl_multitask(self): """Two tasks with two train payloads and two labelsets each""" N = 600 T = 2 tasks = create_tasks(T) model = MetalModel(tasks, verbose=False) payloads = create_payloads(N, T, batch_size=2) metrics_dict = self.trainer.train_model(model, payloads, verbose=False) # For 3 payloads, each of 2 tasks each has 2 label sets self.assertEqual(len(metrics_dict), len(SPLITS) * T**2) for metric, score in metrics_dict.items(): self.assertGreater(score, 0.9)
def get_baseline_scores(tasks, seed, eval_payload): # Baseline Model task_name = tasks[0].name model = MetalModel(tasks, seed=seed, verbose=False) trainer = MultitaskTrainer(seed=seed) trainer.train_model( model, payloads, checkpoint_metric="{}/{}/{}_gold/accuracy".format( task_name, eval_payload.name, task_name), checkpoint_metric_mode="max", checkoint_best=True, writer="tensorboard", optimizer="adamax", lr=5e-5, l2=1e-3, log_every=0.1, score_every=0.1, n_epochs=10, progress_bar=True, checkpoint_best=True, checkpoint_cleanup=False, ) return model.score(eval_payload)
def load_data_and_model(model_path, task_names, split, bert_model="bert-base-uncased"): """ Loads the model specified by model_path and dataset specified by task_name, split. """ # Create DataLoader max_len = 200 dl_kwargs = {"batch_size": 1, "shuffle": False} if not isinstance(task_names, list): task_names = [task_name for task_name in task_names.split(",")] # Load best model for specified task tasks, payloads = create_glue_tasks_payloads( task_names=task_names, bert_model=bert_model, max_len=max_len, dl_kwargs=dl_kwargs, splits=[split], max_datapoints=-1, generate_uids=True, ) # Load and EVAL model model_path = os.path.join(model_path, "best_model.pth") model = MetalModel(tasks, verbose=False, device=0) model.load_weights(model_path) model.eval() for payload in payloads: if payload.split == split: break return model, payload.data_loader
# "COLA", # "SST2", # "MNLI", # "RTE", # "WNLI", # "QQP", # "MRPC", # "STSB", # "QNLI" ], split_prop=0.99, # max_datapoints = 100, dl_kwargs={"batch_size": 2}, ) model = MetalModel(tasks) trainer = MultitaskTrainer(**writer_dict) trainer.train_model( model, payloads, lr=0.00001, l2=0, log_every=1, checkpoint_metric="model/train/loss", checkpoint_metric_mode="min", checkpoint_best=True, n_epochs=3, progress_bar=True, )
def test_mmtl_training(self): model = MetalModel(self.tasks, verbose=False) trainer = MultitaskTrainer(verbose=True) trainer.train_model( model, self.payloads, n_epochs=1, aggregate_metric_fns=[glue_score] )
# Create a valid set from train set task_config["split_prop"] = float(task_config["split_prop"]) # Sampler will be used and handles shuffle automatically task_config["dl_kwargs"]["shuffle"] = False task_config["splits"] = ["train", "test"] else: task_config["splits"] = ["train", "valid", "test"] # Getting primary task names task_names = [task_name for task_name in args.tasks.split(",")] # Getting tasks tasks, payloads = create_glue_tasks_payloads(task_names, **task_config) model_config["verbose"] = False model = MetalModel(tasks, **model_config) # add metadata to trainer_config that will be logged to disk trainer_config["n_paramaters"] = sum(p.numel() for p in model.parameters() if p.requires_grad) if trainer_config["verbose"]: print(f"Task config:\n{task_config}") print(f"Model config:\n{model_config}") print(f"Trainer config:\n{trainer_config}") # Overwrite run_dir to only use one checkpoint dir # if args.run_dir is None: # trainer_config["writer_config"]["run_dir"] = strftime("%Y_%m_%d") if args.run_name is None:
dl_kwargs = {"batch_size": args.batch_size, "shuffle": False} predictions = {} inv_fns = {} for names, model_dirs in models.items(): task_names = names.split(",") for model_dir in model_dirs: with open(os.path.join(model_dir, "task_config.json")) as f: task_config = json.load(f) # TODO: find a nicer way to get task names # create model task_config["splits"] = [] tasks, _ = create_glue_tasks_payloads(task_names, **task_config) model = MetalModel(tasks, verbose=False, device=args.device) if not args.use_task_checkpoints or len(tasks) == 1: # load model weights model_path = os.path.join(model_dir, "best_model.pth") print(f"Loading model checkpoint: {model_path}") model.load_weights(model_path) model.eval() for task in tasks: # only predict for specified tasks if task.name in task_names: # reload task specific checkpoints for MTL models if len(tasks) > 1 and args.use_task_checkpoints: model_path = os.path.join(
foo_input = nn.Linear(2, 10) bar_input = foo_input # nn.Linear(100, 7) foo_head = nn.Linear(10, 2) bar_head = nn.Linear(10, 2) foo_data = make_dataloader(5000) bar_data = make_dataloader(2000) foo = Task("foo_task", {"train": foo_data, "valid": foo_data}, foo_input, foo_head) bar = Task("bar_task", {"train": bar_data, "valid": bar_data}, bar_input, bar_head) # baz = Task("baz_task", "baz_head", [make_dataloader(100), None, None]) tasks = [foo, bar] model = MetalModel(tasks, device=-1, verbose=False) trainer = MultitaskTrainer() trainer.train_model( model, tasks, n_epochs=3, lr=0.1, progress_bar=True, log_every=1, score_every=1, checkpoint_best=True, checkpoint_metric="foo_task/valid/accuracy", checkpoint_metric_mode="max", ) for batch in foo.data_loaders["train"]: