Multitask learning fails with LM1b and SquaAD

[JohannesTK]

Description

Combining multi task learning with Squad fails because of different data structures.

Traceback

"Incompatible return values of true_fn and false_fn: {}".format(e))
ValueError: Incompatible return values of true_fn and false_fn: The two structures don't have the same nested structure.

First structure: type=dict str={'targets': <tf.Tensor 'cond/IteratorGetNext:1' shape=(?,) dtype=int64>, 'batch_prediction_key': <tf.Tensor 'cond/IteratorGetNext:0' shape=(1,) dtype=int64>}

Second structure: type=dict str={'context': <tf.Tensor 'cond/IteratorGetNext_1:1' shape=(?,) dtype=int64>, 'batch_prediction_key': <tf.Tensor 'cond/IteratorGetNext_1:0' shape=(1,) dtype=int64>, 'targets': <tf.Tensor 'cond/IteratorGetNext_1:2' shape=(?,) dtype=int64>}

More specifically: The two dictionaries don't have the same set of keys. First structure has keys type=list str=['targets', 'batch_prediction_key'], while second structure has keys type=list str=['context', 'batch_prediction_key', 'targets']

Reproduce

1. Register problem

# coding=utf-8
"""Data generators for LM1B and SquadConcat combined data-set."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensor2tensor.data_generators import lm1b
from tensor2tensor.data_generators import multi_problem
from tensor2tensor.data_generators import text_problems
from tensor2tensor.utils import registry
from tensor2tensor.data_generators import squad

@registry.register_problem
class LanguagemodelLm1bSquadConcat(multi_problem.MultiProblem):
  """LM1b and SquadConcat mixed problem class for multitask learning."""

  def __init__(self, was_reversed=False, was_copy=False):
    super(LanguagemodelLm1bSquadConcat, self).__init__(was_reversed, was_copy)
    self.task_list.append(lm1b.LanguagemodelLm1b32k())
    self.task_list.append(problem.SquadConcat())

2. Run training

t2t-trainer --t2t_usr_dir=custom_problem \
        --problem=languagemodel_lm1b_squad_concat \
        --data_dir=~/t2t_data \
        --model=transformer \
        --hparams_set=transformer_base \
        --output_dir=~/t2t_train/multitask \
        --train_steps=50000

Environment information

Ubuntu 16.04
python 3.52

tensor2tensor==1.8.0
tensorboard==1.10.0
tensorflow==1.10.0

[2877992943]

same error here,

self.task_list.append(lm1b.LanguagemodelLm1b32k())
self.task_list.append(sequence2sequence_problem)

and

First structure has keys type=list str=['targets', 'batch_prediction_key'], while second structure has keys type=list str=['inputs', 'batch_prediction_key', 'targets'], 

the primary task must be language model and thus the structure of the two task cannot be the same

[zjh-nudger]

any solutions?