Is it possible to script this module?

username: Shisho_Sama
cooked: <p>Hello everyone, hope you are all having a great day.<br>
I was wondering if its possible to convert the following module into torch script:</p>
<pre><code class="lang-python">class PriorBox(torch.nn.Module):
    def __init__(self):
        super(PriorBox, self).__init__()
    
    # @torch.jit.script
    def forward(self, minSizes, steps, clip, image_size):
        anchors = []
        feature_maps = [[ceil(image_size[0]/step), ceil(image_size[1]/step)] for step in steps]

        for k, f in enumerate(feature_maps):
            min_sizes = minSizes[k]
            for i, j in product(range(f[0]), range(f[1])):
                for min_size in min_sizes:
                    s_kx = min_size / image_size[1]
                    s_ky = min_size / image_size[0]
                    dense_cx = [x * steps[k] / image_size[1] for x in [j + 0.5]]
                    dense_cy = [y * steps[k] / image_size[0] for y in [i + 0.5]]
                    for cy, cx in product(dense_cy, dense_cx):
                        anchors += [cx, cy, s_kx, s_ky]

        # back to torch land
        output = torch.Tensor(anchors).view(-1, 4)
        if clip:
            output.clamp_(max=1, min=0)
        return output
</code></pre>
<p>I have reimplemented this in libtorch but its slow as snail!<br>
here it is in case its needed:</p>
<pre><code class="lang-cpp">struct PriorBox : torch::nn::Module
{
    PriorBox(const std::vector&lt;std::pair&lt;torch::Tensor, torch::Tensor&gt;&gt;&amp; min_sizes, const std::vector&lt;int&gt;&amp; steps, const bool&amp; clip, const std::pair&lt;int, int&gt;&amp; image_size, const std::string&amp; phase = "train")
    {
        this-&gt;min_sizes = min_sizes;
        this-&gt;steps = steps;
        this-&gt;clip = clip;
        this-&gt;image_size = image_size;
        this-&gt;phase = phase;
        this-&gt;name = name;
        for (auto&amp; step : this-&gt;steps)
        {
            auto height = torch::tensor(float(this-&gt;image_size.first) / step);
            auto width = torch::tensor(float(this-&gt;image_size.second) / step);
            this-&gt;feature_maps.emplace_back(std::make_pair(torch::ceil(height), torch::ceil(width)));
        }
    }
    
    torch::Tensor forward()
    {
        std::vector&lt;torch::Tensor&gt; anchors;
        int i = -1;
        for (auto&amp; fmap : this-&gt;feature_maps)
        {
            auto min_sizes = this-&gt;min_sizes[++i];
            auto result = torch::cartesian_prod({ torch::arange(c10::Scalar(0), c10::Scalar(fmap.first.item().toInt()), c10::Scalar(1)),
                                                  torch::arange(c10::Scalar(0), c10::Scalar(fmap.second.item().toInt()), c10::Scalar(1)) });
            for (int idx = 0; idx &lt;= result.sizes()[0] - 1; idx++)
            {
                //takes around 0.006 ms
                auto i_ = result[idx][0];
                auto j_ = result[idx][1];
                //takes around 0.20 ms to 0.30 ms
                for (auto&amp; min_size : { min_sizes.first, min_sizes.second })
                {
                    auto s_kx = min_size / float(this-&gt;image_size.second);
                    auto s_ky = min_size / float(this-&gt;image_size.first);
                    // takes around 0.037 ms
                    torch::Tensor dense_cx = (j_ + 0.5) * this-&gt;steps[i] / this-&gt;image_size.second;
                    torch::Tensor dense_cy = (i_ + 0.5) * this-&gt;steps[i] / this-&gt;image_size.first;
                    //takes around 0.02ms
                    auto result_cy_cx = torch::cartesian_prod({ dense_cy.unsqueeze(0), dense_cx.unsqueeze(0) });
                    //this takes around 0.010ms
                    for (int l = 0; l &lt;= result_cy_cx.sizes()[0] - 1; l++)
                    {
                        auto cy = result_cy_cx[l][0].unsqueeze(0);
                        auto cx = result_cy_cx[l][1].unsqueeze(0);
                        anchors.emplace_back(torch::cat({ cx, cy, s_kx, s_ky }));
                    }
                }
            }
        }
        //takes around 5ms!
        auto output = torch::stack(anchors).view({ -1,4 });
        //takes around 0 ms!
        if (this-&gt;clip)
            output.clamp_(0, 1);
        
        return output;
    }

    std::vector&lt;std::pair&lt;torch::Tensor, torch::Tensor&gt;&gt; min_sizes;
    std::vector&lt;int&gt; steps;
    bool clip;
    std::pair&lt;int, int&gt; image_size;
    std::string phase;
    std::string name;
    std::vector&lt;std::pair&lt;torch::Tensor, torch::Tensor&gt;&gt; feature_maps;
};

</code></pre>
<p>So I thought maybe converting this into a Torchscript and loading it in C++ would be a better idea.<br>
Since we are dealing with loops, etc we cant simply jit trace the module. So I guess I’m stuck with scripting.<br>
When I tried to do scripting I kept getting different errors such asrange cannot be used as a value:</p>
<pre><code class="lang-auto">RuntimeError:
range cannot be used as a value:
  File "P:\ligen\layers\functions\prior_box.py", line 19
        for k, f in enumerate(feature_maps):
            min_sizes = minSizes[k]
            for i, j in product(range(int(f[0])), range(int(f[1]))):
                                ~~~~~~~~~~~~~~ &lt;--- HERE
                for min_size in min_sizes:
                    s_kx = min_size / image_size[1]
</code></pre>
<p>So my question is, considering this, is it portable to torchscript? if so What am I missing here? how should I go about this?<br>
If not, what are my other options?<br>
Thanks a lot in advance</p>

username: SplitInfinity
cooked: <p>In general, scripting requires you to modify your input program so that it only uses the features that <code>torch.jit.script</code> supports. Take a look at the <a href="https://pytorch.org/docs/stable/jit_language_reference.html" rel="noopener nofollow ugc">TorchScript Language Reference</a> for more information.</p>

username: Shisho_Sama
cooked: <p>Thanks, I apprecaite your kind help, but I know about it.<br>
The issue is I’m not sure how I should be going about this! how can I substitute the range as I can’t find it in the supported list and I’m at a complete loss of idea on how to go about this!</p>

username: Prashant
cooked: <p>Torchscript allows for range. I think its just your product(range, range) that’s the problem. Maybe just nest for loops.</p>

username: Shisho_Sama
cooked: <p>Thanks a lot.<br>
I did try the range outside of the product and it failed with the same error nevertheless.<br>
I also substituted the product with the PyTorch’s built-in cartesian-prod if I remember correctly but that seems not to work either.</p>

username: vferrer
cooked: <p>Torchscript code won’t be much faster. Torchscript doesn’t speed up your code by a factor of 2. If<br>
If you reimplemented in libtorch and code is slow, also it will be with Torchscript.</p>
<p>Finally, I think you could avoid some list comprehension like with dense_cx and dense_cy</p>

username: Shisho_Sama
cooked: <p>Thanks a lot. I already stripped down much of libtorch machinery and resorted back to the old plain for loops and got around 3x speed up!<br>
Having that part jit scripted should in theory at least help in performance, as I believe (and expect) the instructions are optimized as well by the engine) at least this is what I hoped for.<br>
I guess I’ll be able to hopefully extract a bit more performance out of it (making it around 4x faster), but thats still no way around the performance I get in Python!<br>
I’m not sure if this is because something is wrong with libtorch’s <code>cartesian_prod</code> implementation that takes too much compared to its counterpart (Python’s product) or some thing else that I’m missing completely. I know the implementation posted above is not efficient by any means, but I really didn’t expect this to be this slow!<br>
To give you an impression, Python code runs around 20,30 ms while this takes around 350!450 ms!<br>
(after my last changes, it now runs at 70~100ms )but still way slower than Python!</p>