[pytorch-1.0rc1] using torch.jit.trace() to trace a network that takes 2 input

username: tlm
cooked: <p>Hi,<br>
I am trying to trace using torch.jit.trace() a network that takes 2 tensors (z and x) as input and produces one tensor as output (y).</p>
<p>I can successfully trace models with 1 input and 1 output as follows:</p>
<pre><code class="lang-auto">  sample_input = torch.rand(1,3,256,256)
  traced_module_feature = torch.jit.trace(model, sample_input)
  traced_module_feature.save("traced_model.pt")
</code></pre>
<p>However, I have no clue on how to trace models that take 2 inputs. I tried the following which produces error as given below:</p>
<pre><code class="lang-auto">z = torch.rand(1,3,127,127)
x = torch.rand(1,3,256,256)
traced_module_rpn = torch.jit.trace(model, { 'zf':z, 'xf':x } )
</code></pre>
<p>and the output is</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "export_to_cpp.py", line 57, in &lt;module&gt;
    main()
  File "export_to_cpp.py", line 52, in main
    traced_module_rpn = torch.jit.trace(model.rpn_model, { 'zf':z, 'xf':x } )
  File "/home/tlm/anaconda3/envs/svt2_pyth1/lib/python3.7/site-packages/torch/jit/__init__.py", line 565, in trace
    module._create_method_from_trace('forward', func, example_inputs)
RuntimeError: Only tensors and (possibly nested) tuples of tensors are supported as inputs or outputs of traced functions (toIValue at /opt/conda/conda-bld/pytorch-nightly_1538562647654/work/torch/csrc/jit/pybind_utils.h:74)
frame #0: &lt;unknown function&gt; + 0x3fe53f (0x7f98c5fb353f in /home/tlm/anaconda3/envs/svt2_pyth1/lib/python3.7/site-packages/torch/_C.cpython-37m-x86_64-linux-gnu.so)
frame #1: &lt;unknown function&gt; + 0x463a2b (0x7f98c6018a2b in /home/tlm/anaconda3/envs/svt2_pyth1/lib/python3.7/site-packages/torch/_C.cpython-37m-x86_64-linux-gnu.so)
frame #2: &lt;unknown function&gt; + 0x1a665d (0x7f98c5d5b65d in /home/tlm/anaconda3/envs/svt2_pyth1/lib/python3.7/site-packages/torch/_C.cpython-37m-x86_64-linux-gnu.so)
&lt;omitting python frames&gt;
frame #19: __libc_start_main + 0xf0 (0x7f98da43b830 in /lib/x86_64-linux-gnu/libc.so.6)
</code></pre>
<p>I am using pytorch-1.0-rc1 with cuda-9.0 and python 3.7.</p>
<p>Thank you</p>

username: tlm
cooked: <p>and, this is the network that I am trying to trace:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F


class RPN(nn.Module):
    def __init__(self):
        super(RPN, self).__init__()

    def forward(self, z_f, x_f):
        raise NotImplementedError


def conv2d_group(x, kernel):
    batch = kernel.size()[0]
    pk = kernel.view(-1, x.size()[1], kernel.size()[2], kernel.size()[3])
    px = x.view(1, -1, x.size()[2], x.size()[3])
    po = F.conv2d(px, pk, groups=batch)
    po = po.view(batch, -1, po.size()[2], po.size()[3])
    return po


class UPChannelRPN(RPN):
    def __init__(self, anchor_num=5, feature_in=256, feature_out=256):
        super(UPChannelRPN, self).__init__()

        self.anchor_num = anchor_num
        self.feature_in = feature_in
        self.feature_out = feature_out

        self.cls_output = 2 * self.anchor_num
        self.loc_output = 4 * self.anchor_num

        self.template_cls_conv = nn.Conv2d(self.feature_in, self.feature_out * self.cls_output, kernel_size=3)
        self.template_loc_conv = nn.Conv2d(self.feature_in, self.feature_out * self.loc_output, kernel_size=3)

        self.search_cls_conv = nn.Conv2d(self.feature_in, self.feature_out, kernel_size=3)
        self.search_loc_conv = nn.Conv2d(self.feature_in, self.feature_out, kernel_size=3)

        self.loc_adjust = nn.Conv2d(self.loc_output, self.loc_output, kernel_size=1)

    def forward(self, z_f, x_f):
        cls_kernel = self.template_cls_conv(z_f)
        loc_kernel = self.template_loc_conv(z_f)

        cls_feature = self.search_cls_conv(x_f)
        loc_feature = self.search_loc_conv(x_f)

        pred_cls = conv2d_group(cls_feature, cls_kernel)
        pred_loc = self.loc_adjust(conv2d_group(loc_feature, loc_kernel))
        return pred_cls, pred_loc
</code></pre>
<p>I am trying to trace <code>UPChannelRPN()</code>. This network does not include any conditionals (to the best of my knowledge) and therefore I think it should be possible to trace this network using <code>torch.jit.trace</code></p>

username: tlm
cooked: <p>The trace works when I use:</p>
<pre><code class="lang-auto">traced_module_rpn = torch.jit.trace(model.rpn_model, [ Variable(z), Variable(x) ] )
</code></pre>
<p>but with this warning:</p>
<pre><code class="lang-auto">rpn.py:18: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  po = F.conv2d(px, pk, groups=batch)
</code></pre>
<p>where <code>line 18: of rpn.py</code> is</p>
<pre><code class="lang-auto">    po = F.conv2d(px, pk, groups=batch)
</code></pre>
<p>Any help in understanding this warning message is highly appreciated.</p>

username: JeeLee
cooked: <aside class="quote no-group" data-username="tlm" data-post="3" data-topic="26527">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/0ea827/48.png" class="avatar"> tlm:</div>
<blockquote>
<p>traced_module_rpn = torch.jit.trace(model.rpn_model, [ Variable(z), Variable(x) ] )</p>
</blockquote>
</aside>
<p>I think you can try tuple type</p>

username: Nikos_Mparoutis
cooked: <p>Hello, i would like to see the trace graphically through<br>
torch.jit.trace:<br>
traced_net = torch.jit.trace(net,inputs, optimize=True, check_trace=True, check_inputs=None, check_tolerance=1e-05)<br>
print(traced_net)<br>
make_dot_from_trace(traced_net)    &lt;—</p>
<p>—&gt;AttributeError: ‘TopLevelTracedModule’ object has no attribute ‘set_graph’<br>
If there is not graph then what ?</p>

username: Agarg
cooked: <p>Any insights on how to solve this. Similar problem that needs to be solved. trying to trace the model with a RPN head with 2 inputs.</p>

username: vriez
cooked: <p>This works</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(1, 1, 3)

    def forward(self, x, y):
        return self.conv(x)

    def weighted_kernel_sum(self, weight):
        return weight * self.conv.weight


n = Net()

example_weight = torch.rand(1, 1, 3, 3)
example_forward_input = torch.rand(1, 1, 3, 3)

# targets
H, W = 1000, 700
num_classes = 10
image_id = torch.as_tensor(1693029438, dtype=torch.int32)
num_objects = 1200
image = torch.randn((H, W))
masks = torch.randint(0, 1, (num_objects, H, W), dtype=torch.uint8) + 255
areas = torch.randint(100, 20000, (1, num_objects), dtype=torch.int64)
boxes = torch.randint(100, H * W, (num_objects, 4), dtype=torch.int64)
labels = torch.randint(1, num_classes, (1, num_objects), dtype=torch.int64)
iscrowd = torch.zeros(num_objects, dtype=torch.int64)

target = {}
target["image_id"] = image_id
target["boxes"] = boxes
target["labels"] = labels
target["area"] = areas
target["iscrowd"] = iscrowd
target["masks"] = masks


# Trace a specific method and construct `ScriptModule` with
# a single `forward` method
module = torch.jit.trace(n.forward, (example_forward_input, target))

# Trace a module (implicitly traces `forward`) and construct a
# `ScriptModule` with a single `forward` method
module = torch.jit.trace(n, (example_forward_input, target))

# Trace specific methods on a module (specified in `inputs`), constructs
# a `ScriptModule` with `forward` and `weighted_kernel_sum` methods
inputs = {'forward' : (example_forward_input, target), 'weighted_kernel_sum' : example_weight}
module = torch.jit.trace_module(n, inputs)
module.save("jit.pt")

</code></pre>
<p>This was taken from <a href="https://pytorch.org/docs/stable/generated/torch.jit.trace_module.html" rel="noopener nofollow ugc">here</a></p>