Multitask model prediction accuracy lower after converting to C++

marchss: <p>I have a multitask model which uses ModuleDict. Its <code>forward(self, a, a_len, b, b_len, task_id)</code> has argument task_id.</p>
<p>After converting the pytorch model checkpoint into C++ by tracing, the C++ prediction accuracy is 5% lower than pytorch prediction accuracy.</p>
<p>If I train each task individually, i.e., into a separate model and checkpoint file, then convert each of them into C++, then the accuracy matches.</p>

tom: <p>Tracing only covers the specific code-path executed during tracing, all other paths and control-flow information cannot be recorded<br>
You could trace the various submodules and <a href="https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting" rel="nofollow noopener">combine them using scripting</a> to get an single exportable model encompassing everything.</p>
<p>Best regards</p>
<p>Thomas</p>

marchss: <p><a class="mention" href="/u/tom">@tom</a> Thank you so much. I haven’t fully understand your answer yet. The link is nice, however, I don’t know how to apply it into <code>for-loops</code>, i.e. iterate all possible <code>task_id</code> as input. Would you please provide some example?</p>
<p>Below is the code snippet of the Model class. It contains  <code>ModuleDict</code> whose key is <code>task_id</code>. The <code>forward</code> function has <code>task_id</code> as argument.</p>
<p>This is a control flow, similar to <code>if-else</code>, since  <code>ModuleDict</code> internally uses <code>for-loops</code>?</p>
<p>So, when tracing the model. Shall we trace ALL <code>task_id</code> or tracing only one task_id (variable channel in code) is enough as shown below?</p>
<pre><code class="lang-auto">channel = torch.ones(1, dtype=torch.int64)
traced_script_module = torch.jit.trace(model, (premise, premise_length, hypotheses, hypotheses_length, channel))

output = traced_script_module(premise, premise_length, hypotheses, hypotheses_length, channel)
traced_script_module.save('deploy-trace-multitask.pt')
</code></pre>
<p>Code snippet for Model class’s definition</p>
<pre><code class="lang-auto">        self._word_embedding = nn.Embedding(self.vocab_size,
                                            self.embedding_dim,
                                            padding_idx=padding_idx,
                                            _weight=embeddings)

        if self.dropout:
            self._rnn_dropout = RNNDropout(p=self.dropout) #shared by all tasks
            # self._rnn_dropout = nn.Dropout(p=self.dropout)

        self._encoding = Seq2SeqEncoder(nn.LSTM,
                                        self.embedding_dim,
                                        self.hidden_size,
                                        bidirectional=True)

        #multi-task
        self._attention = nn.ModuleDict({})
        self._projection = nn.ModuleDict({})
        self._classification = nn.ModuleDict({})
        for channel in channels_list:
            self.update(channel)

        # Initialize all weights and biases in the model.
        self.apply(_init_esim_weights)

    def update(self, channel):
        channel = str(channel)
        self._attention.update({channel : SoftmaxAttention()})

        self._projection.update({channel : nn.Sequential(nn.Linear(4*2*self.hidden_size, self.hidden_size), nn.ReLU())})

        self._classification.update({channel : nn.Sequential(nn.Dropout(p=self.dropout),
                                             nn.Linear(4*self.hidden_size,
                                                       self.hidden_size),
                                             nn.Tanh(),
                                             nn.Dropout(p=self.dropout),
                                             nn.Linear(self.hidden_size,
                                                       self.num_classes))})

    def forward(self,
                premises,
                premises_lengths,
                hypotheses,
                hypotheses_lengths,
                channel_tensor): #must be a tensor
        """
        Args:
            premises: A batch of varaible length sequences of word indices
                representing premises. The batch is assumed to be of size
                (batch, premises_length).
            premises_lengths: A 1D tensor containing the lengths of the
                premises in 'premises'.
            hypothesis: A batch of varaible length sequences of word indices
                representing hypotheses. The batch is assumed to be of size
                (batch, hypotheses_length).
            hypotheses_lengths: A 1D tensor containing the lengths of the
                hypotheses in 'hypotheses'.

        Returns:
            logits: A tensor of size (batch, num_classes) containing the
                logits for each output class of the model.
            probabilities: A tensor of size (batch, num_classes) containing
                the probabilities of each output class in the model.
        """
        channel_id = channel_tensor.item()
        channel = str(channel_id)
        premises_mask = get_mask(premises, premises_lengths).to(self.device)
        hypotheses_mask = get_mask(hypotheses, hypotheses_lengths)\
            .to(self.device)

        embedded_premises = self._word_embedding(premises)
        embedded_hypotheses = self._word_embedding(hypotheses)

        if self.dropout:
            embedded_premises = self._rnn_dropout(embedded_premises)
            embedded_hypotheses = self._rnn_dropout(embedded_hypotheses)

        encoded_premises = self._encoding(embedded_premises,
                                          premises_lengths)
        encoded_hypotheses = self._encoding(embedded_hypotheses,
                                            hypotheses_lengths)

        attended_premises, attended_hypotheses =\
            self._attention[channel](encoded_premises, premises_mask,
                            encoded_hypotheses, hypotheses_mask)
        """ rest of the code are omitted """
</code></pre>

tom: <p>You need to trace the tasks separately and then write the bit combining them in torch.script, calling the traced modules. The link has a small recipe for that.</p>
<p>Best regards</p>
<p>Thomas</p>

marchss: <p><a class="mention" href="/u/tom">@tom</a> Thank you for your patience, sir. Would you please clarify the step of “combining them”?</p>
<p>I indeed checked the <code>for-loop</code> example in the link which accumulates result in each for-loop iteration.</p>
<p>However, here, each channel’s result is final, and excludes each other. When using multitask model, it expose argument for each individual channel. For example, passing in channel “games” will produce a final score for itself (no need to combine with other channels), while passing in channel “fashion” will produce another final score for itself.</p>
<p>What shall be returned in <code>loop_in_traced_fn</code>? Returning any channel’s result  will discard other channels’.</p>
<pre><code class="lang-auto">@torch.jit.script
def loop_in_traced_fn(premise, premise_length, hypotheses, hypotheses_length, channel_list):
	for channel in channel_list:
		result = model(premise, premise_length, hypotheses, hypotheses_length, channel)
		#what shall be returned here? since each channel's result exclude each other. 

#channel_list = ["games", "fashion", "news", "food"]
traced = torch.jit.trace(loop_in_traced_fn, premise, premise_length, hypotheses, hypotheses_length, channel_list)
</code></pre>

tom: <p>You can accumulate them in a list and return the list <code>mylist = []</code> and  <code>mylist.append(...)</code> should work… (Or add them, <code>torch.cat</code>, whatever…).</p>
<p>Best regards</p>
<p>Thomas</p>

marchss: <p>Thank you, Thomas.</p>
<ol>
<li>
<p>would you please point me to the exact example you are referring to? I tried several but no luck yet. If we write a wrapper<code>class Composition</code>,  trace each subtask separately and combine them by calling <code>forward(self, inputs)</code> whose <code>inputs</code> is <code>["games", "fashion", "news", "food"]</code>. This means, the final saved trace C++ model won’t be able to accept single channel like <code>games</code>, right?</p>
</li>
<li>
<p>For my code in previous post,<br>
It seems that <code>torch.jit.trace(loop_in_traced_fn, inputs)</code> can ONLY accepts the <code>inputs</code> as <code>tuple</code>, so that  <code>loop_in_traced_fn</code> will get the tensor inside the <code>inputs</code> tuple. Am I correct?</p>
</li>
</ol>
<p>It still not working yet, with error message shown below. I have also tried all kinds of input such as list of tuples, list of tensors, etc.</p>
<pre><code class="lang-auto">
RuntimeError: 
Tensor cannot be used as a tuple:
@torch.jit.script
def loop_in_traced_fn(channel_tuple):
    result = []
    for channel in channel_tuple:
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~...  &lt;--- HERE
        print(channel)
        channel = torch.tensor([channel], dtype=torch.int64)
        output = model(premise, premise_length, hypotheses, hypotheses_length, channel)
        result.append(output)
    return result
</code></pre>

tom: <p>Could you do the following, please?</p>
<ul>
<li>Trace the individual models as you did before. Store them in a dictionary with string keys.</li>
<li>Write as simple a Python method that you can think of that does what you want the  final, combined model to do, using only the dictionary with the traced models (and any inputs).</li>
</ul>
<p>Best regards</p>
<p>Thomas</p>

marchss: <p><a class="mention" href="/u/tom">@tom</a> Thank you very much, Sir. Yes, I followed what you said. Below is the error message. What am I missing?</p>
<pre><code class="lang-auto">TypeError: 'dict' object for attribute 'module_dict' is not a valid constant.
Valid constants are:
  1. a nn.ModuleList
  2. a value of type {bool, float, int, str, NoneType, function, device, layout, dtype}
  3. a list or tuple of (2)
</code></pre>
<p>Below is a snippet of combining them</p>
<pre><code class="lang-auto">#trace each channel's model separately
channel_tensors = []
module_dict = dict()
for channel_id in channels_list:
    channel = torch.tensor([channel_id], dtype=torch.int64)
    channel_tensors.append(channel)
    traced_script_module = torch.jit.trace(model, (premise, premise_length, hypotheses, hypotheses_length, channel))
    module_dict[channel_id] = traced_script_module

#combine each channel's model together
class MyScriptModule(torch.jit.ScriptModule):
    __constants__ = ['module_dict']

    def __init__(self, module_dict):
        super(MyScriptModule, self).__init__()
        self.module_dict = module_dict

    @torch.jit.script_method
    def forward(self,
                premises,
                premises_lengths,
                hypotheses,
                hypotheses_lengths,
                channel): #channel must be a tensor
        channel_id = channel.item()
        return self.module_dict[channel_id](premise, premise_length, hypotheses, hypotheses_length, channel)

my_script_module = MyScriptModule(module_dict)
my_script_module.save("deploy-trace-multitask.channel_all.pt")
</code></pre>

marchss: <p><a class="mention" href="/u/tom">@tom</a> I also tried another aproach by defining a simple python function to trace, and get errors as below.</p>
<pre><code class="lang-auto">    _jit_script_compile(mod, ast, _rcb, get_default_args(fn))
RuntimeError: 
python value of type 'dict' cannot be used as a value:
@torch.jit.script
def loop_in_traced_fn(premise, premise_length, hypotheses, hypotheses_length, channel):
    channel_id = channel.item()
    result = module_dict[channel_id](premise, premise_length, hypotheses, hypotheses_length, channel)
             ~~~~~~~~~~~ &lt;--- HERE
    return result
</code></pre>
<p>Below is the code snipeet</p>
<pre><code class="lang-auto">#trace each channel's model separately
channel_tensors = []
module_dict = dict()
for channel_id in channels_list:
    channel = torch.tensor([channel_id], dtype=torch.int64)
    channel_tensors.append(channel)
    traced_script_module = torch.jit.trace(model, (premise, premise_length, hypotheses, hypotheses_length, channel))
    module_dict[channel_id] = traced_script_module

@torch.jit.script
def loop_in_traced_fn(premise, premise_length, hypotheses, hypotheses_length, channel):
    channel_id = channel.item()
    result = module_dict[channel_id](premise, premise_length, hypotheses, hypotheses_length, channel)
    return result

channel = channel_tensors[0]
traced = torch.jit.trace(loop_in_traced_fn, (premise, premise_length, hypotheses, hypotheses_length, channel))
traced.save("deploy-trace-multitask.channel_all.pt")
</code></pre>

marchss: <aside class="quote no-group" data-username="tom" data-post="8" data-topic="51011">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>Write as simple a Python method that you can think of that does what you want the final, combined model to do, using only the dictionary with the traced models (and any inputs).</p>
</blockquote>
</aside>
<p>Python method in both approaches report error message like “python dict type is NOT supported in Script”. The dict is used to map each channel into its traced model. This is the confusing part.</p>
<p>I guess this python method also requires to be traced and saved as the big wrapper model for all channel’s individually traced models.</p>

tom: <p>It turns out that it’s a bit more trick, I’m still working on it…</p>

marchss: <p>Thank you so much for your kind help. The main issue is how to <strong>trace and save</strong> the python method  which combines all channels’ traced module.</p>

marchss: <p>Hi, Thomas. Is it possible? Thanks.</p>

marchss: <p>Anyone knows, please?</p>