Pytorch c++ extensions versus Tensorflow op

Ahmed_Abdelaziz: <p>I converted a TensorFlow based model into PyTorch. The tf code used c++ op that I converted into PyTorch C++ Extention. I compared the timings of both running a model for 100 epochs, the tf op takes less than 3 seconds to finish for a 32 batch sizes while for PyTorch it takes on average 4 minutes on CPU and on GPU it takes double the time for the same batch size. Please see below the attached code sample from PyTorch</p>
<pre><code class="lang-auto">torch::Tensor extract_spans(
    torch::Tensor span_scores,
    torch::Tensor candidate_starts,
    torch::Tensor candidate_ends,
    torch::Tensor num_output_spans,
    int max_sentence_length,
    bool _sort_spans,
    bool _suppress_crossing
) {

    int num_sentences = span_scores.size(0);
    int num_input_spans = span_scores.size(1);
    int max_num_output_spans = 0;


    for (int i = 0; i &lt; num_sentences; i++) {

      if (num_output_spans[i].item&lt;int64_t&gt;() &gt; max_num_output_spans) {
        max_num_output_spans = num_output_spans[i].item&lt;int64_t&gt;();
      }
    }


    std::vector&lt;std::vector&lt;int&gt;&gt; sorted_input_span_indices(num_sentences,
                                                            std::vector&lt;int&gt;(num_input_spans));

    torch::Tensor output_span_indices = torch::ones({num_sentences, max_num_output_spans});

    for (int i = 0; i &lt; num_sentences; i++) {
      std::iota(sorted_input_span_indices[i].begin(), sorted_input_span_indices[i].end(), 0);
      std::sort(sorted_input_span_indices[i].begin(), sorted_input_span_indices[i].end(),
                [&amp;span_scores, &amp;i](int j1, int j2) {
                 if (j1 &gt;= span_scores.size(1) || j1 &lt; 0 || j2 &gt;= span_scores.size(1) || j2 &lt; 0) {
                    return false;
                 }

                  return span_scores[i][j2].item&lt;int64_t&gt;() &lt; span_scores[i][j1].item&lt;int64_t&gt;();
                });
    }


    for (int l = 0; l &lt; num_sentences; l++) {
      std::vector&lt;int&gt; top_span_indices;
      std::unordered_map&lt;int, int&gt; end_to_earliest_start;
      std::unordered_map&lt;int, int&gt; start_to_latest_end;
      int current_span_index = 0, num_selected_spans = 0;

      while (num_selected_spans &lt; num_output_spans[l].item&lt;int64_t&gt;() &amp;&amp; current_span_index &lt; num_input_spans) {
        int i = sorted_input_span_indices[l][current_span_index];
        bool any_crossing = false;
        if (_suppress_crossing) {
          const int&amp; start = candidate_starts[l][i].item&lt;int64_t&gt;();
          const int&amp; end = candidate_ends[l][i].item&lt;int64_t&gt;();

          for (int j = start; j &lt;= end; ++j) {
            if (j &gt; start) {
              auto latest_end_iter = start_to_latest_end.find(j);
              if (latest_end_iter != start_to_latest_end.end() &amp;&amp; latest_end_iter-&gt;second &gt; end) {
                // Given (), exists [], such that ( [ ) ]
                any_crossing = true;
                break;
              }
            }
            if (j &lt; end) {
              auto earliest_start_iter = end_to_earliest_start.find(j);
              if (earliest_start_iter != end_to_earliest_start.end() &amp;&amp; earliest_start_iter-&gt;second &lt; start) {
                // Given (), exists [], such that [ ( ] )
                any_crossing = true;
                break;
              }
            }
          }
        }
        if (!any_crossing) {
          if (_sort_spans) {
            top_span_indices.push_back(i);
          } else {
            output_span_indices[l][num_selected_spans] = i;
          }
          ++num_selected_spans;
          if (_suppress_crossing) {
            // Update data struct.
            const int&amp; start = candidate_starts[l][i].item&lt;int64_t&gt;();
            const int&amp; end = candidate_ends[l][i].item&lt;int64_t&gt;();
            auto latest_end_iter = start_to_latest_end.find(start);
            if (latest_end_iter == start_to_latest_end.end() || end &gt; latest_end_iter-&gt;second) {
              start_to_latest_end[start] = end;
            }
            auto earliest_start_iter = end_to_earliest_start.find(end);
            if (earliest_start_iter == end_to_earliest_start.end() || start &lt; earliest_start_iter-&gt;second) {
              end_to_earliest_start[end] = start;
            }
          }
        }
        ++current_span_index;
      }
      // Sort and produce span indices.
      if (_sort_spans) {
        std::sort(top_span_indices.begin(), top_span_indices.end(),
                [&amp;candidate_starts, &amp;candidate_ends, &amp;l] (int i1, int i2) {
                 if (i1 &gt;= candidate_starts.size(1) || i1 &lt; 0 || i2 &gt;= candidate_starts.size(1) || i2 &lt; 0) {
                    return false;
                 }
                  if (candidate_starts[l][i1].item&lt;int64_t&gt;() &lt; candidate_starts[l][i2].item&lt;int64_t&gt;()) {
                    return true;
                  } else if (candidate_starts[l][i1].item&lt;int64_t&gt;() &gt; candidate_starts[l][i2].item&lt;int64_t&gt;()) {
                    return false;
                  } else if (candidate_ends[l][i1].item&lt;int64_t&gt;() &lt; candidate_ends[l][i2].item&lt;int64_t&gt;()) {
                    return true;
                  } else if (candidate_ends[l][i1].item&lt;int64_t&gt;() &gt; candidate_ends[l][i2].item&lt;int64_t&gt;()) {
                    return false;
                  } else {
                    return i1 &lt; i2;
                  }
                });

        for (int i = 0; i &lt; num_output_spans[l].item&lt;int64_t&gt;(); ++i) {
          output_span_indices[l][i] = top_span_indices[i];
        }
      }

      // Pad with the last selected span index to ensure monotonicity.
      int last_selected = num_selected_spans - 1;
      if (last_selected &gt;= 0) {
        for (int i = num_selected_spans; i &lt; max_num_output_spans; ++i) {
          output_span_indices[l][i]= output_span_indices[l][last_selected].item&lt;int64_t&gt;();
        }
      }

    }



    return output_span_indices;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("extract_spans", &amp;extract_spans, "extract_spans");
}
</code></pre>
<p>The equivalent in tf is</p>
<pre><code class="lang-auto">REGISTER_OP("ExtractSpans")
.Input("span_scores: float32")
.Input("candidate_starts: int32")
.Input("candidate_ends: int32")
.Input("num_output_spans: int32")
.Input("max_sentence_length: int32")
.Attr("sort_spans: bool")
.Attr("suppress_crossing: bool")
.Output("output_span_indices: int32");

class ExtractSpansOp : public OpKernel {
public:
  explicit ExtractSpansOp(OpKernelConstruction* context) : OpKernel(context) {
    OP_REQUIRES_OK(context, context-&gt;GetAttr("sort_spans", &amp;_sort_spans));
    OP_REQUIRES_OK(context, context-&gt;GetAttr("suppress_crossing", &amp;_suppress_crossing));
  }

  void Compute(OpKernelContext* context) override {
    TTypes&lt;float&gt;::ConstMatrix span_scores = context-&gt;input(0).matrix&lt;float&gt;();
    TTypes&lt;int32&gt;::ConstMatrix candidate_starts = context-&gt;input(1).matrix&lt;int32&gt;();
    TTypes&lt;int32&gt;::ConstMatrix candidate_ends = context-&gt;input(2).matrix&lt;int32&gt;();
    TTypes&lt;int32&gt;::ConstVec num_output_spans = context-&gt;input(3).vec&lt;int32&gt;();
    int max_sentence_length = context-&gt;input(4).scalar&lt;int32&gt;()();

    int num_sentences = span_scores.dimension(0);
    int num_input_spans = span_scores.dimension(1);
    int max_num_output_spans = 0;
    for (int i = 0; i &lt; num_sentences; i++) {
      if (num_output_spans(i) &gt; max_num_output_spans) {
        max_num_output_spans = num_output_spans(i);
      }
    }

    Tensor* output_span_indices_tensor = nullptr;
    TensorShape output_span_indices_shape({num_sentences, max_num_output_spans});
    OP_REQUIRES_OK(context, context-&gt;allocate_output(0, output_span_indices_shape,
                                                     &amp;output_span_indices_tensor));
    TTypes&lt;int32&gt;::Matrix output_span_indices = output_span_indices_tensor-&gt;matrix&lt;int32&gt;();

    std::vector&lt;std::vector&lt;int&gt;&gt; sorted_input_span_indices(num_sentences,
                                                            std::vector&lt;int&gt;(num_input_spans));
    for (int i = 0; i &lt; num_sentences; i++) {
      std::iota(sorted_input_span_indices[i].begin(), sorted_input_span_indices[i].end(), 0);
      std::sort(sorted_input_span_indices[i].begin(), sorted_input_span_indices[i].end(),
                [&amp;span_scores, &amp;i](int j1, int j2) {
                  return span_scores(i, j2) &lt; span_scores(i, j1);
                });
    }

    for (int l = 0; l &lt; num_sentences; l++) {
      std::vector&lt;int&gt; top_span_indices;
      std::unordered_map&lt;int, int&gt; end_to_earliest_start;
      std::unordered_map&lt;int, int&gt; start_to_latest_end;
      int current_span_index = 0,
          num_selected_spans = 0;
      while (num_selected_spans &lt; num_output_spans(l) &amp;&amp; current_span_index &lt; num_input_spans) {
        int i = sorted_input_span_indices[l][current_span_index];
        bool any_crossing = false;
        if (_suppress_crossing) {
          const int&amp; start = candidate_starts(l, i);
          const int&amp; end = candidate_ends(l, i);
          for (int j = start; j &lt;= end; ++j) {
            if (j &gt; start) {
              auto latest_end_iter = start_to_latest_end.find(j);
              if (latest_end_iter != start_to_latest_end.end() &amp;&amp; latest_end_iter-&gt;second &gt; end) {
                // Given (), exists [], such that ( [ ) ]
                any_crossing = true;
                break;
              }
            }
            if (j &lt; end) {
              auto earliest_start_iter = end_to_earliest_start.find(j);
              if (earliest_start_iter != end_to_earliest_start.end() &amp;&amp; earliest_start_iter-&gt;second &lt; start) {
                // Given (), exists [], such that [ ( ] )
                any_crossing = true;
                break;
              }
            }
          }
        }
        if (!any_crossing) {
          if (_sort_spans) {
            top_span_indices.push_back(i);
          } else {
            output_span_indices(l, num_selected_spans) = i;
          }
          ++num_selected_spans;
          if (_suppress_crossing) {
            // Update data struct.
            const int&amp; start = candidate_starts(l, i);
            const int&amp; end = candidate_ends(l, i);
            auto latest_end_iter = start_to_latest_end.find(start);
            if (latest_end_iter == start_to_latest_end.end() || end &gt; latest_end_iter-&gt;second) {
              start_to_latest_end[start] = end;
            }
            auto earliest_start_iter = end_to_earliest_start.find(end);
            if (earliest_start_iter == end_to_earliest_start.end() || start &lt; earliest_start_iter-&gt;second) {
              end_to_earliest_start[end] = start;
            }
          }
        }
        ++current_span_index;
      }
      // Sort and produce span indices.
      if (_sort_spans) {
        std::sort(top_span_indices.begin(), top_span_indices.end(),
                [&amp;candidate_starts, &amp;candidate_ends, &amp;l] (int i1, int i2) {
                  if (candidate_starts(l, i1) &lt; candidate_starts(l, i2)) {
                    return true;
                  } else if (candidate_starts(l, i1) &gt; candidate_starts(l, i2)) {
                    return false;
                  } else if (candidate_ends(l, i1) &lt; candidate_ends(l, i2)) {
                    return true;
                  } else if (candidate_ends(l, i1) &gt; candidate_ends(l, i2)) {
                    return false;
                  } else {
                    return i1 &lt; i2;
                  }
                });
        for (int i = 0; i &lt; num_output_spans(l); ++i) {
          output_span_indices(l, i) = top_span_indices[i];
        }
      }
      // Pad with the last selected span index to ensure monotonicity.
      int last_selected = num_selected_spans - 1;
      if (last_selected &gt;= 0) {
        for (int i = num_selected_spans; i &lt; max_num_output_spans; ++i) {
          output_span_indices(l, i) = output_span_indices(l, last_selected);
        }
      }
    }
  }
private:
  bool _sort_spans, _suppress_crossing;
};

REGISTER_KERNEL_BUILDER(Name("ExtractSpans").Device(DEVICE_CPU), ExtractSpansOp);

</code></pre>
<p>So why pytorch is much slower than tf in that case? and why running on GPU takes more time compared to running the same code on a cpu?</p>

ptrblck: <p>How did you measure the timings and did you synchronize both codes?<br>
I’m not familiar enough with TensorFlow and don’t know how the code can be synchronized, but in PyTorch you should call <code>torch.cuda.synchronize()</code> before starting and stopping the timer. Otherwise your measurements are invalid.</p>

Ahmed_Abdelaziz: <p>Thanks for your reply. I at least logged the time before and after inside the c++ code assuming it’s a blocking code. I now converted the code back to python and it works very fast compared to the c++ setup, I don’t know why.</p>