1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
|
#include <torch/csrc/jit/mobile/train/sequential.h>
#include <torch/types.h>
#include <algorithm>
#include <cstddef>
#include <vector>
namespace torch::jit::mobile {
SequentialSampler::SequentialSampler(size_t size) : size_(size) {}
void SequentialSampler::reset(std::optional<size_t> new_size) {
if (new_size.has_value()) {
size_ = *new_size;
}
index_ = 0;
}
std::optional<std::vector<size_t>> SequentialSampler::next(size_t batch_size) {
const auto remaining_indices = size_ - index_;
if (remaining_indices == 0) {
return std::nullopt;
}
std::vector<size_t> index_batch(std::min(batch_size, remaining_indices));
for (auto& i : index_batch) {
i = index_++;
}
return index_batch;
}
void SequentialSampler::save(serialize::OutputArchive& archive) const {
TORCH_CHECK(
false, "Serialization of SequentialSampler not supported on mobile.");
}
void SequentialSampler::load(serialize::InputArchive& archive) {
TORCH_CHECK(
false, "Serialization of SequentialSampler not supported on mobile.");
}
size_t SequentialSampler::index() const noexcept {
return index_;
}
} // namespace torch::jit::mobile
|
