Custom dataset using libtorch review request

username: yngtodd
cooked: <p>Hey everyone,</p>
<p>I am running into a bit of trouble with an undefined reference when creating a custom dataset class using <code>libtorch</code>. May I ask for a code review to help clarify some things?</p>
<p>here is my <code>data.hpp</code>:</p>
<pre><code class="lang-auto">#pragma once

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

namespace rock {
namespace data {
namespace datasets {

/// Random dataset.
class RandomDataset : public torch::data::Dataset&lt;RandomDataset&gt; {
     public:
          /// The mode in which the dataset is loaded.
          enum class Mode { kTrain, kTest };

          explicit RandomDataset(Mode mode = Mode::kTrain);

	      /// Returns the `Example` at the given `index`.
	     torch::data::Example&lt;&gt; get(size_t index) override;

	     /// Returns the size of the dataset.
	     torch::optional&lt;size_t&gt; size() const override;

	      /// Returns true if this is the training subset of MNIST.
          bool is_train() const noexcept;

          /// Returns all images stacked into a single tensor.
          const torch::Tensor&amp; images() const;

          /// Returns all targets stacked into a single tensor.
          const torch::Tensor&amp; labels() const;

     private:
	      torch::Tensor images_, labels_;
};
} // namespace datasets
} // namespace data
} // namespace rock
</code></pre>
<p>And my <code>data.cpp</code>:</p>
<pre><code class="lang-auto">#include &lt;torch/torch.h&gt;
#include "data.hpp"

namespace rock {
namespace data {
namespace datasets {
namespace {

constexpr uint32_t kTrainSize = 60000;
constexpr uint32_t kNumChannels = 3;
constexpr uint32_t kImageRows = 256;
constexpr uint32_t kImageColumns = 256;

// Create random images
torch::Tensor load_images() {
    torch::Tensor tensor = torch::randint(
        /*low=*/0, /*high=*/255, 
        {kTrainSize, kImageRows, kImageColumns, kNumChannels}
    ); 
    
    return tensor.to(torch::kFloat32).div_(255);
}

// Create labels - all ones.
torch::Tensor load_labels() {
    torch::Tensor tensor = torch::ones({kTrainSize}, torch::kInt);
    return tensor.to(torch::kInt64);
}
} // namespace

RandomDataset::RandomDataset(Mode mode)
    : images_(load_images()),
      labels_(load_labels()) {}

torch::data::Example&lt;&gt; RandomDataset::get(size_t index) {
    return {images_[index], labels_[index]};
}

torch::optional&lt;size_t&gt; RandomDataset::size() const {
    return images_.size(0);
}

bool RandomDataset::is_train() const noexcept {
    return images_.size(0) == kTrainSize;
}

const torch::Tensor&amp; RandomDataset::images() const {
    return images_;
}

const torch::Tensor&amp; RandomDataset::labels() const {
    return labels_;
}

} // namespace datasets
} // namespace data
} // namespace rock
</code></pre>
<p>I saw that <a class="mention" href="/u/mhubii">@mhubii</a> linked a good reference implementation for a custom dataset here: <a href="https://discuss.pytorch.org/t/libtorch-how-to-use-torch-datasets-for-custom-dataset/34221" class="inline-onebox">(libtorch) How to use torch::data::datasets for custom dataset?</a> . Not quite sure where I am getting stuck on this reference.</p>
<p>I was structuring my code to be similar to the <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/api/src/data/datasets/mnist.cpp" rel="nofollow noopener">mnist.cpp</a>.</p>
<p>I appreciate your help and comments!</p>

username: yngtodd
cooked: <p>Messed up my <code>CMakeLists.txt</code> <img src="https://discuss.pytorch.org/images/emoji/apple/man_facepalming.png?v=9" title=":man_facepalming:" class="emoji" alt=":man_facepalming:">. That took care of the reference. Everything compiles fine, but now seems to hang when initializing that dataset:</p>
<pre><code class="lang-auto">auto dataset = rock::data::datasets::RandomDataset()
    .map(torch::data::transforms::Stack&lt;&gt;());
</code></pre>

username: yngtodd
cooked: <p>And another false alarm, I was creating random tensors of shape <code>{60_000, 256, 256, 3}</code> and my workstation was taking longer than expected. Everything is running well!</p>