def test_multitask_custom_attachments(self): """Attach the task heads at user-specified layers""" edges = [(0, 1)] cards = [2, 2] tg = TaskHierarchy(cards, edges) em = MTEndModel( layer_out_dims=[2, 8, 4], task_graph=tg, seed=1, verbose=False, dropout=0.0, task_head_layers=[1, 2], ) self.assertEqual(em.task_map[1][0], 0) self.assertEqual(em.task_map[2][0], 1) em.train( self.Xs[0], self.Ys[0], self.Xs[1], self.Ys[1], verbose=False, n_epochs=10, ) score = em.score(self.Xs[2], self.Ys[2], reduce="mean", verbose=False) self.assertGreater(score, 0.95)
def test_multitask_top(self): """Attach all task heads to the top layer""" edges = [] cards = [2, 2] tg = TaskGraph(cards, edges) em = MTEndModel( layer_out_dims=[2, 8, 4], task_graph=tg, seed=1, verbose=False, dropout=0.0, task_head_layers="top", ) top_layer = len(em.config["layer_out_dims"]) - 1 self.assertEqual(len(em.task_map[top_layer]), em.t) em.train( self.Xs[0], self.Ys[0], self.Xs[1], self.Ys[1], verbose=False, n_epochs=10, ) score = em.score(self.Xs[2], self.Ys[2], reduce="mean", verbose=False) self.assertGreater(score, 0.95)
def test_multitask_two_modules(self): """Accept a different representation for each task""" edges = [] cards = [2, 2] tg = TaskGraph(cards, edges) em = MTEndModel( layer_out_dims=[2, 8, 4], task_graph=tg, seed=1, verbose=False, dropout=0.0, input_modules=[IdentityModule(), IdentityModule()], task_head_layers="top", ) Xs = [] for i, X in enumerate(self.Xs): Xs.append([X[:, 0], X[:, 1]]) em.train( Xs[0], self.Ys[0], Xs[1], self.Ys[1], verbose=False, n_epochs=10 ) score = em.score(Xs[2], self.Ys[2], reduce="mean", verbose=False) self.assertGreater(score, 0.95)
def test_multitask_custom(self): edges = [(0, 1)] cards = [2, 2] tg = TaskHierarchy(edges, cards) em = MTEndModel( task_graph=tg, seed=1, verbose=False, dropout=0.0, layer_output_dims=[2, 8, 4], task_head_layers=[1, 2], ) em.train( self.Xs[0], self.Ys[0], self.Xs[1], self.Ys[1], verbose=False, n_epochs=10, ) score = em.score(self.Xs[2], self.Ys[2], reduce='mean', verbose=False) self.assertGreater(score, 0.95)
def test_multitask_custom_heads(self): """Accept a different representation for each task""" edges = [] cards = [2, 2] tg = TaskGraph(cards, edges) em = MTEndModel( layer_out_dims=[2, 8, 4], task_graph=tg, seed=1, verbose=False, dropout=0.0, head_modules=[nn.Linear(8, 2), nn.Linear(4, 2)], task_head_layers=[1, 2], ) em.train( self.Xs[0], self.Ys[0], self.Xs[1], self.Ys[1], verbose=False, n_epochs=10, ) score = em.score(self.Xs[2], self.Ys[2], reduce="mean", verbose=False) self.assertGreater(score, 0.95)