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_model(
(Xs[0], self.Ys[0]),
dev_data=(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_custom_modules(self):
"""Test custom input/head modules"""
input_module = nn.Sequential(IdentityModule(), nn.Linear(2, 10))
middle_modules = [nn.Linear(10, 8), IdentityModule()]
head_module = nn.Sequential(nn.Linear(8, 2), IdentityModule())
em = EndModel(
seed=1,
input_module=input_module,
middle_modules=middle_modules,
head_module=head_module,
layer_out_dims=[10, 8, 8],
verbose=False,
)
Xs, Ys = self.single_problem
em.train(
Xs[0],
Ys[0],
Xs[1],
Ys[1],
n_epochs=5,
verbose=False,
show_plots=False,
)
score = em.score(Xs[2], Ys[2], verbose=False)
self.assertGreater(score, 0.95)
def _build_input_layer(self, input_module):
if input_module is None:
input_module = IdentityModule()
output_dim = self.config["layer_out_dims"][0]
input_layer = self._make_layer(
input_module,
"input",
self.config["input_layer_config"],
output_dim=output_dim,
)
return input_layer
def _build_input_layer(self, input_modules):
if input_modules is None:
output_dim = self.config["layer_out_dims"][0]
input_modules = IdentityModule()
if isinstance(input_modules, list):
input_layer = [self._make_layer(mod) for mod in input_modules]
else:
input_layer = self._make_layer(input_modules, output_dim)
return input_layer