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"]: