Returning more than 2 tensors in libtorch

a_d: <p>Hello all ,<br>
I am writing a simple dataloader for an object detection task and trying to return 3 tensors from the <code>get</code> method. However I intended to return a tuple containing all three but I am unable to do so as the return type is <code>torch::data::Example</code>. Could anyone please guide me as to how I can return all three tensors?<br>
TIA</p>

ptrblck: <p>Wouldn’t returning all 3 tensors work with <code>torch::data::Example</code> as given e.g. in the <code>get</code> method of your dataset:</p>
<pre><code class="lang-auto">torch::data::Example&lt;&gt; get(size_t i) {
    return {torch::ones(i), torch::ones(i), torch::ones(i)};
  }
</code></pre>

a_d: <p>I did try that but it threw an error saying no matching constructor found and I was only able to return at most two tensors.<br>
Anyway I will try that once again. Thanks a lot for replying ptrbk.</p>

krzim: <p>I just stumbled across this exact same issue when trying to create a custom dataset that returns 3 tensors vs. 2. I’m not crazy well versed in C++ but it looks like <a class="mention" href="/u/ptrblck">@ptrblck</a>’s solution wouldn’t work (and didn’t for me at least) according to <a href="https://github.com/pytorch/pytorch/blob/484534f3be4409f8ae759f9968710b7f4afa5093/torch/csrc/api/include/torch/data/example.h" rel="noopener nofollow ugc">the source code</a>. I think we’re limited to 1 or 2 outputs based on the templates in that header. I’d be interested to hear of any potential workarounds or solutions to this issue.</p>

ptrblck: <p>Did you try to write a custom <code>Example</code> class with more than two return types?<br>
The linked class shows:</p>
<pre><code class="lang-auto">template &lt;typename Data = Tensor, typename Target = Tensor&gt;
</code></pre>
<p>so I would assume you can reuse this template and add more types to it?</p>

krzim: <p>Wow that was a quick response! I’m actually working on that approach now. I’ll see how it goes and report back with what I find.</p>
<p>I’m mostly a Python dev though so I’m a bit out of my element here.</p>
<p>EDIT: It looks like writing a custom <code>Example</code> class and providing it as a template argument when creating a custom dataset subclass (per <a href="https://github.com/pytorch/pytorch/blob/6cef59487ab2e8a7124924f9a76cab381c638ac5/torch/csrc/api/include/torch/data/datasets/base.h#L75" rel="noopener nofollow ugc">this line here</a>) should work. I’ll write up a minimal example here when I figure it out.</p>

krzim: <p>Apologies in advance if this doesn’t compile directly on a copy/paste. I wrote this based loosely on my actual implementation but haven’t compiled or tested this exact bit of code. I also  haven’t fully tested this approach in a training loop but the approach compiles and can be used to do batched inference with my TorchScript model. Note that I had to create a custom <code>Example</code> and custom <code>Stack</code> class to accommodate the extra model input in transforms as well.</p>
<p>I’m not great with C++ but I’m wondering if there’s some way that parameter packs/variadic templating can be used to remove the 2 input limit on custom datasets? I started probing that with the <code>Example</code> class but don’t have any strong ideas on how to carry that forward into the transforms or anything.</p>
<p>It looks like <code>example.data</code> and <code>example.target</code> are used quite often in examples, transforms, etc. vs. taking the indexing/unpacking approach to getting inputs from a batch in Python. Maybe it’s worth exploring the Python approach here as well.</p>
<pre><code class="lang-cpp">#include &lt;tuple&gt;
#include &lt;vector&gt;
#include &lt;fliesystem&gt;

#include &lt;torch/torch.h&gt;

template &lt;typename Data = torch::Tensor, typename Target = torch::Tensor,
          typename Mask = torch::Tensor&gt;
struct Example {
  using DataType = Data;
  using TargetType = Target;
  using MaskType = Mask;

  Data data;
  Target target;
  Mask mask;

  Example() = default;
  Example(Data data, Target target, Mask mask)
      : data(std::move(data)), target(std::move(target)),
        mask(std::move(mask)) {}
}

template &lt;typename ExampleType&gt;
struct Stack : public torch::data::transforms::Collation&lt;ExampleType&gt; {
  ExampleType apply_batch(std::vector&lt;ExampleType&gt; examples) override {
    std::vector&lt;torch::Tensor&gt; xs, ys, masks;
    xs.reserve(examples.size());
    ys.reserve(examples.size());
    masks.reserve(examples.size());
    for (auto &amp;example : examples) {
      xs.push_back(std::move(example.data));
      ys.push_back(std::move(example.target));
      masks.push_back(std::move(std::get&lt;0&gt;(example.rest)));
    }
    return {torch::stack(xs), torch::stack(ys), torch::stack(masks)};
  }
};

using MyExample = Example&lt;torch::Tensor, torch::Tensor, torch::Tensor&gt;;
class MyDataset : public torch::data::datasets::Dataset&lt;MyDataset, MyExample&gt; {
private:
  std::vector&lt;torch::Tensor&gt; xs, ys, masks;

public:
  MyDataset(const std::filesystem::path dataFpath) {
    xs, ys, masks = loadDataFromFile(dataFpath);
  }
  virtual ~MyDataset();
  MyExample get(size_t index) override {
    torch::Tensor x = xs[index];
    torch::Tensor y = ys[index];
    torch::Tensor mask = masks[index];

    return {x, y, mask};
  }
  torch::optional&lt;size_t&gt; size() const override { return xs.size(); }
}

auto main() -&gt; int {
  const std::filesystem::path dataDir = "/path/to/data/";
  const int batchSize = 32;
  auto dataset = MyDataset(dataDir).map(Stack&lt;MyExample&gt;());
  auto dataLoader =
      torch::data::make_data_loader&lt;torch::data::samplers::SequentialSampler&gt;(
          std::move(dataset), batchSize);
}
</code></pre>

ptrblck: <p>Thanks for sharing your approach!<br>
I’m sure the <code>libtorch</code> interface could be polished a bit more, but I also think that training in C++/<code>libtorch</code> is still an edge case (compared to inference) which is why the training API might not have gotten a lot attention.</p>

krzim: <p>I agree on both fronts. There aren’t many usecases where training in C++ makes more sense than training in Python and converting a model to TorchScript for inference in C++.</p>
<p>I’m working on model training in C++ for an iOS app ultimately. Unfortunately I’ll need to be able to do on device training for my usecase so I’m certain I’ll come across a lot of unpolished edges as I move forward. I’ll be sure to open issues and make PRs where I can along the way.</p>
<p>Thanks for the help on this!</p>

osar: <p>Hello krzim,<br>
I am facing a similar problem, where my network has two branches, each awaiting for a different size of tensor. I could not succeed in making your code example work. Didi you end up with a solution ? Would you mind sharing it ?<br>
For my problem, I was also wondering if my Example could contain a pair of tensors as the data and another tensor as the target. But, here too, I could not obtain in a working code. I am not easy enough with C++ for that. Any idea welcome.<br>
Thanks</p>

osar: <p>Just to share my solution, I ended up with this simpler code:</p>
<pre><code class="lang-auto">auto makeDataPrototypes()
{
  std::vector&lt;std::tuple&lt;torch::Tensor, torch::Tensor, torch::Tensor&gt;&gt; dataset;
  for (int i = 0; i &lt; 100; i++)
    dataset.push_back(std::make_tuple(torch::randn({2, 41, 41}), torch::randn({2}), torch::randn({1})));
  return dataset;
}

struct Example {
  torch::Tensor input1;
  torch::Tensor input2;
  torch::Tensor output;

  Example() = default;
  Example(torch::Tensor input1, torch::Tensor input2, torch::Tensor output)
      : input1(std::move(input1)), input2(std::move(input2)), output(std::move(output)) {}
};

struct Stack : public torch::data::transforms::Collation&lt;Example&gt; {
  Example apply_batch(std::vector&lt;Example&gt; examples) override {
    std::vector&lt;torch::Tensor&gt; inputs1, inputs2, outputs;
    inputs1.reserve(examples.size());
    inputs2.reserve(examples.size());
    outputs.reserve(examples.size());
    for (auto &amp;example : examples) {
      inputs1.push_back(std::move(example.input1));
      inputs2.push_back(std::move(example.input2));
      outputs.push_back(std::move(example.output));
    }
    return {torch::stack(inputs1), torch::stack(inputs2), torch::stack(outputs)};
  }
};

class MyDataset : public torch::data::datasets::Dataset&lt;MyDataset, Example&gt; {

  private:
    std::vector&lt;std::tuple&lt;torch::Tensor, torch::Tensor, torch::Tensor&gt;&gt; dataset_;
  
  public:
    MyDataset() :
      dataset_(makeDataPrototypes()) {
        std::cout &lt;&lt; "nb proto = " &lt;&lt; dataset_.size() &lt;&lt; '\n';
    }
    Example get(size_t index) override {
      auto tuple = dataset_[index];
      return {std::get&lt;0&gt;(tuple), std::get&lt;1&gt;(tuple), std::get&lt;2&gt;(tuple)};
    }
    torch::optional&lt;size_t&gt; size() const override { 
      return dataset_.size(); 
    }
};

int main(int argc, char *argv[]) {
  
  auto dataset = MyDataset().map(Stack());
  std::cout &lt;&lt; "nb proto in main = " &lt;&lt; dataset.size().value() &lt;&lt; '\n';
  auto loader = torch::data::make_data_loader&lt;torch::data::samplers::RandomSampler&gt;(std::move(dataset), 10);
  std::cout &lt;&lt; "nb proto after loader creation = " &lt;&lt; dataset.size().value() &lt;&lt; '\n';
 /* for (auto&amp; batch : *loader) 
    auto outputs = net-&gt;forward(batch.input1, batch.input2);*/
}
</code></pre>
<p>The only thing that is still mysterious for me is that nb proto in main after loader creation is 0:</p>
<pre><code class="lang-auto">nb proto = 100
nb proto in main = 100
nb proto after loader creation = 0
</code></pre>
<p>It is due to the std::move, but is it required? Anyway the loop on batch examples seems to work fine.</p>

MaxC2: <p>With regard to the implementation of <code>torch::Tensor</code> (use of an <code>intrusive_ptr</code>), <code>std::move</code> is not necessary because it breaks with the externality of the reference counter. (the semantic move is used internally in the base class)</p>
<p>As stated in the <code>TensorBody.h</code> documentation</p>
<blockquote>
<p>Tensor is a “generic” object holding a pointer to the underlying <code>TensorImpl</code> object, which has an embedded reference count. In this way, Tensor is similar to <code>boost::intrusive_ptr</code>.</p>
<p>For example:</p>
<pre><code class="lang-auto">void func(Tensor a) {
   Tensor b = a;
   ...
}
</code></pre>
<p>In this example, when we say <code>Tensor b = a</code>, we are creating a new object that points to the same underlying TensorImpl, and bumps its reference count. When b goes out of scope, the destructor decrements the reference count by calling release() on the TensorImpl it points to. The existing constructors, operator overloads, etc. take care to implement the correct semantics.</p>
</blockquote>
<p>It is therefore possible that by using std::move in your <code>Example</code> constructor, you lose the underlaying construction of your Tensor’s.</p>
<p>Not sure at 100%, I do not use the move semantic with the libtorch API.</p>

gnscc: <p>If you’re finding that the other options aren’t working for you, you might consider defining a custom return type for your dataset. This method seems more aligned with the intended usage of LibTorch’s dataset functionalities.</p>
<p>Firstly, create a structure for your custom return type and then use this type with <code>torch::data::datasets::Dataset</code>. In this context, the first template parameter is the type of your custom dataset class (<code>MyCustomDataset</code>), and the second parameter is the type of the data item your dataset will return (<code>MyCustomExample</code>). By default, the data item type for a dataset in LibTorch is <code>torch::data::Dataset::Example&lt;at::Tensor, at::Tensor&gt;</code>, but you can customize it as shown below:</p>
<pre data-code-wrap="cpp"><code class="lang-cpp">struct MyCustomExample {
    torch::Tensor image;
    torch::Tensor mask;
    std::string classname;
};

class MyCustomDataset : public torch::data::datasets::Dataset&lt;MyCustomDataset, MyCustomExample&gt; {
  public:
    // ...
    MyCustomExample get(size_t index) override;
    // ...
};
</code></pre>