def test_save_load(self):
fd, checkpoint_path = tempfile.mkstemp()
task1 = create_task("task1")
task2 = create_task("task2")
# Make task2's second linear layer have different weights
task2.module_pool["linear2"] = nn.Linear(2, 2)
model = MultitaskModel([task1])
self.assertTrue(
torch.eq(
task1.module_pool["linear2"].weight,
model.module_pool["linear2"].module.weight,
).all())
model.save(checkpoint_path)
model = MultitaskModel([task2])
self.assertFalse(
torch.eq(
task1.module_pool["linear2"].weight,
model.module_pool["linear2"].module.weight,
).all())
model.load(checkpoint_path)
self.assertTrue(
torch.eq(
task1.module_pool["linear2"].weight,
model.module_pool["linear2"].module.weight,
).all())
os.close(fd)
def test_bad_tasks(self):
with self.assertRaisesRegex(ValueError, "Found duplicate task"):
MultitaskModel(tasks=[self.task1, self.task1])
with self.assertRaisesRegex(ValueError, "Unrecognized task type"):
MultitaskModel(tasks=[self.task1, {"fake_task": 42}])
with self.assertRaisesRegex(ValueError, "Unsuccessful operation"):
task1 = create_task("task1")
task1.op_sequence[0].inputs[0] = (0, 0)
model = MultitaskModel(tasks=[task1])
X_dict = self.dataloader.dataset.X_dict
model.forward(X_dict, [task1.name])
def test_partially_empty_batch(self):
dataset = create_dataloader("task1", shuffle=False).dataset
dataset.Y_dict["task1"][0] = -1
model = MultitaskModel([self.task1])
loss_dict, count_dict = model.calculate_loss(dataset.X_dict,
dataset.Y_dict)
self.assertEqual(count_dict["task1"], 9)
def test_empty_batch(self):
dataset = create_dataloader("task1", shuffle=False).dataset
dataset.Y_dict["task1"] = torch.full_like(dataset.Y_dict["task1"], -1)
model = MultitaskModel([self.task1])
loss_dict, count_dict = model.calculate_loss(dataset.X_dict,
dataset.Y_dict)
self.assertFalse(loss_dict)
self.assertFalse(count_dict)
def test_twotask_partial_overlap_model(self):
"""Add two tasks with overlapping modules and flows"""
task1 = create_task("task1", module_suffixes=["A", "A"])
task2 = create_task("task2", module_suffixes=["A", "B"])
model = MultitaskModel(tasks=[task1, task2])
self.assertEqual(len(model.task_names), 2)
self.assertEqual(len(model.op_sequences), 2)
self.assertEqual(len(model.module_pool), 3)
def test_load_on_cleanup(self) -> None:
log_manager_config = {"counter_unit": "epochs", "evaluation_freq": 1}
checkpointer = Checkpointer(**log_manager_config, checkpoint_dir=self.test_dir)
log_manager = LogManager(
n_batches_per_epoch=2, checkpointer=checkpointer, log_writer=None
)
classifier = MultitaskModel([])
best_classifier = log_manager.cleanup(classifier)
self.assertEqual(best_classifier, classifier)
def test_score(self):
model = MultitaskModel([self.task1])
metrics = model.score([self.dataloader])
# deterministic random tie breaking alternates predicted labels
self.assertEqual(metrics["task1/dataset/train/accuracy"], 0.4)
# test dataframe format
metrics_df = model.score([self.dataloader], as_dataframe=True)
self.assertTrue(isinstance(metrics_df, pd.DataFrame))
self.assertEqual(metrics_df.at[0, "score"], 0.4)
def test_no_input_spec(self):
# Confirm model doesn't break when a module does not specify specific inputs
dataset = create_dataloader("task", shuffle=False).dataset
task = Task(
name="task",
module_pool=nn.ModuleDict({"identity": nn.Identity()}),
op_sequence=[Operation("identity", [])],
)
model = MultitaskModel(tasks=[task], dataparallel=False)
outputs = model.forward(dataset.X_dict, ["task"])
self.assertIn("_input_", outputs)
def test_checkpointer_load_best(self) -> None:
checkpoint_dir = os.path.join(self.test_dir, "clear")
checkpointer = Checkpointer(
**log_manager_config,
checkpoint_dir=checkpoint_dir,
checkpoint_metric="task/dataset/valid/f1:max",
)
model = MultitaskModel([])
checkpointer.checkpoint(1, model, {"task/dataset/valid/f1": 0.8})
load_model = checkpointer.load_best_model(model)
self.assertEqual(model, load_model)
def test_score_shuffled(self):
# Test scoring with a shuffled dataset
class SimpleVoter(nn.Module):
def forward(self, x):
"""Set class 0 to -1 if x and 1 otherwise"""
mask = x % 2 == 0
out = torch.zeros(x.shape[0], 2)
out[mask, 0] = 1 # class 0
out[~mask, 1] = 1 # class 1
return out
# Create model
task_name = "VotingTask"
module_name = "simple_voter"
module_pool = nn.ModuleDict({module_name: SimpleVoter()})
op0 = Operation(module_name=module_name,
inputs=[("_input_", "data")],
name="op0")
op_sequence = [op0]
task = Task(name=task_name,
module_pool=module_pool,
op_sequence=op_sequence)
model = MultitaskModel([task])
# Create dataset
y_list = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
x_list = [i for i in range(len(y_list))]
Y = torch.LongTensor(y_list * 100)
X = torch.FloatTensor(x_list * 100)
dataset = DictDataset(name="dataset",
split="train",
X_dict={"data": X},
Y_dict={task_name: Y})
# Create dataloaders
dataloader = DictDataLoader(dataset, batch_size=2, shuffle=False)
scores = model.score([dataloader])
self.assertEqual(scores["VotingTask/dataset/train/accuracy"], 0.6)
dataloader_shuffled = DictDataLoader(dataset,
batch_size=2,
shuffle=True)
scores_shuffled = model.score([dataloader_shuffled])
self.assertEqual(scores_shuffled["VotingTask/dataset/train/accuracy"],
0.6)
def test_checkpointer_clear(self) -> None:
checkpoint_dir = os.path.join(self.test_dir, "clear")
checkpointer = Checkpointer(
**log_manager_config,
checkpoint_dir=checkpoint_dir,
checkpoint_metric="task/dataset/valid/f1:max",
checkpoint_clear=True,
)
model = MultitaskModel([])
checkpointer.checkpoint(1, model, {"task/dataset/valid/f1": 0.8})
expected_files = [
"checkpoint_1.pth", "best_model_task_dataset_valid_f1.pth"
]
self.assertEqual(set(os.listdir(checkpoint_dir)), set(expected_files))
checkpointer.clear()
expected_files = ["best_model_task_dataset_valid_f1.pth"]
self.assertEqual(os.listdir(checkpoint_dir), expected_files)
def test_checkpointer_min(self) -> None:
checkpointer = Checkpointer(
**log_manager_config,
checkpoint_dir=self.test_dir,
checkpoint_runway=3,
checkpoint_metric="task/dataset/valid/f1:min",
)
model = MultitaskModel([])
checkpointer.checkpoint(3, model, {
"task/dataset/valid/f1": 0.8,
"task/dataset/valid/f2": 0.5
})
self.assertEqual(
checkpointer.best_metric_dict["task/dataset/valid/f1"], 0.8)
checkpointer.checkpoint(4, model, {"task/dataset/valid/f1": 0.7})
self.assertEqual(
checkpointer.best_metric_dict["task/dataset/valid/f1"], 0.7)
def test_predict(self):
model = MultitaskModel([self.task1])
results = model.predict(self.dataloader)
self.assertEqual(sorted(list(results.keys())), ["golds", "probs"])
np.testing.assert_array_equal(
results["golds"]["task1"],
self.dataloader.dataset.Y_dict["task1"].numpy())
np.testing.assert_array_equal(results["probs"]["task1"],
np.ones((NUM_EXAMPLES, 2)) * 0.5)
results = model.predict(self.dataloader, return_preds=True)
self.assertEqual(sorted(list(results.keys())),
["golds", "preds", "probs"])
# deterministic random tie breaking alternates predicted labels
np.testing.assert_array_equal(
results["preds"]["task1"],
np.array([0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0]),
)
def test_remapped_labels(self):
# Test additional label keys in the Y_dict
# Without remapping, model should ignore them
task_name = self.task1.name
X = torch.FloatTensor([[i, i] for i in range(NUM_EXAMPLES)])
Y = torch.ones(NUM_EXAMPLES).long()
Y_dict = {task_name: Y, "other_task": Y}
dataset = DictDataset(name="dataset",
split="train",
X_dict={"data": X},
Y_dict=Y_dict)
dataloader = DictDataLoader(dataset, batch_size=BATCH_SIZE)
model = MultitaskModel([self.task1])
loss_dict, count_dict = model.calculate_loss(dataset.X_dict,
dataset.Y_dict)
self.assertIn("task1", loss_dict)
# Test setting without remapping
results = model.predict(dataloader)
self.assertIn("task1", results["golds"])
self.assertNotIn("other_task", results["golds"])
scores = model.score([dataloader])
self.assertIn("task1/dataset/train/accuracy", scores)
self.assertNotIn("other_task/dataset/train/accuracy", scores)
# Test remapped labelsets
results = model.predict(dataloader,
remap_labels={"other_task": task_name})
self.assertIn("task1", results["golds"])
self.assertIn("other_task", results["golds"])
results = model.score([dataloader],
remap_labels={"other_task": task_name})
self.assertIn("task1/dataset/train/accuracy", results)
self.assertIn("other_task/dataset/train/accuracy", results)
loss_func=F.mse_loss,
output_func=identity_fn,
scorer=Scorer(metrics=["mse"]))
# -
# ## Model
# With our tasks defined, constructing a model is simple: we simply pass the list of tasks in and the model constructs itself using information from the task flows.
#
# Note that the model uses the names of modules (not the modules themselves) to determine whether two modules specified by separate tasks are the same module (and should share weights) or different modules (with separate weights).
# So because both the `class_task` and `rgb_task` include "base_net" in their module pools, this module will be shared between the two tasks.
# +
from cerbero.models import MultitaskModel
model = MultitaskModel([class_task, rgb_task])
# -
# ### Train Model
# Once the model is constructed, we can train it as we would a single-task model, using the `fit` method of a `Trainer` object. The `Trainer` supports multiple schedules or patterns for sampling from different dataloaders; the default is to randomly sample from them proportional to the number of batches, such that all data points will be seen exactly once before any are seen twice.
# +
from cerbero.trainer import Trainer
trainer_config = {
"progress_bar": True,
"n_epochs": 15,
"lr": 2e-3,
"checkpointing": True
}
def test_no_data_parallel(self):
model = MultitaskModel(tasks=[self.task1, self.task2],
dataparallel=False)
self.assertEqual(len(model.task_names), 2)
self.assertIsInstance(model.module_pool["linear1A"], nn.Module)
op_sequence = [op1, op2]
task = Task(name=task_name,
module_pool=module_pool,
op_sequence=op_sequence)
return task
dataloaders = [create_dataloader(task_name) for task_name in TASK_NAMES]
tasks = [
create_task(TASK_NAMES[0], module_suffixes=["A", "A"]),
create_task(TASK_NAMES[1], module_suffixes=["A", "B"]),
]
model = MultitaskModel([tasks[0]])
class TrainerTest(unittest.TestCase):
def test_trainer_onetask(self):
"""Train a single-task model"""
trainer = Trainer(**base_config)
trainer.fit(model, [dataloaders[0]])
def test_trainer_twotask(self):
"""Train a model with overlapping modules and flows"""
multitask_model = MultitaskModel(tasks)
trainer = Trainer(**base_config)
trainer.fit(multitask_model, dataloaders)
def test_trainer_errors(self):
def test_trainer_twotask(self):
"""Train a model with overlapping modules and flows"""
multitask_model = MultitaskModel(tasks)
trainer = Trainer(**base_config)
trainer.fit(multitask_model, dataloaders)
def test_twotask_none_overlap_model(self):
"""Add two tasks with totally separate modules and flows"""
model = MultitaskModel(tasks=[self.task1, self.task2])
self.assertEqual(len(model.task_names), 2)
self.assertEqual(len(model.op_sequences), 2)
self.assertEqual(len(model.module_pool), 4)
def test_onetask_model(self):
model = MultitaskModel(tasks=[self.task1])
self.assertEqual(len(model.task_names), 1)
self.assertEqual(len(model.op_sequences), 1)
self.assertEqual(len(model.module_pool), 2)