reid from https://github.com/michuanhaohao/reid-strong-baseline
zhangmeng
2020-01-10 c3765bd24fe73747688a0ec2a550f219c9acb384
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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#pragma once
 
#include <torch/data/example.h>
#include <torch/types.h>
 
#include <c10/util/ArrayRef.h>
 
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace torch {
namespace data {
namespace datasets {
template <typename S, typename T>
class MapDataset;
template <typename D, typename T>
MapDataset<D, T> map(D, T); // NOLINT
} // namespace datasets
} // namespace data
} // namespace torch
 
namespace torch {
namespace data {
namespace datasets {
namespace detail {
template <typename T>
struct is_optional : std::false_type {};
template <typename T>
struct is_optional<optional<T>> : std::true_type {};
} // namespace detail
 
/// A dataset that can yield data only in batches.
template <
    typename Self,
    typename Batch = std::vector<Example<>>,
    typename BatchRequest = ArrayRef<size_t>>
class BatchDataset {
 public:
  using SelfType = Self;
  using BatchType = Batch;
  using BatchRequestType = BatchRequest;
  constexpr static bool is_stateful = detail::is_optional<BatchType>::value;
 
  virtual ~BatchDataset() = default;
 
  /// Returns a batch of data given an index.
  virtual Batch get_batch(BatchRequest request) = 0;
 
  /// Returns the size of the dataset, or an empty optional if it is unsized.
  virtual optional<size_t> size() const = 0;
 
  /// Creates a `MapDataset` that applies the given `transform` to this dataset.
  template <typename TransformType>
  MapDataset<Self, TransformType> map(TransformType transform) & {
    return datasets::map(static_cast<Self&>(*this), std::move(transform));
  }
 
  /// Creates a `MapDataset` that applies the given `transform` to this dataset.
  template <typename TransformType>
  MapDataset<Self, TransformType> map(TransformType transform) && {
    return datasets::map(
        std::move(static_cast<Self&>(*this)), std::move(transform));
  }
};
 
/// A dataset that can yield data in batches, or as individual examples.
///
/// A `Dataset` is a `BatchDataset`, because it supports random access and
/// therefore batched access is implemented (by default) by calling the random
/// access indexing function for each index in the requested batch of indices.
/// This can be customized.
template <typename Self, typename SingleExample = Example<>>
class Dataset : public BatchDataset<Self, std::vector<SingleExample>> {
 public:
  using ExampleType = SingleExample;
 
  /// Returns the example at the given index.
  virtual ExampleType get(size_t index) = 0;
 
  /// Returns a batch of data.
  /// The default implementation calls `get()` for every requested index
  /// in the batch.
  std::vector<ExampleType> get_batch(ArrayRef<size_t> indices) override {
    std::vector<ExampleType> batch;
    batch.reserve(indices.size());
    for (const auto i : indices) {
      batch.push_back(get(i));
    }
    return batch;
  }
};
 
/// A `StreamDataset` reprsents a dataset that is a potentially infinite stream.
/// It takes as batch index only a number, which is the batch size, and yields
/// that many elements from the stream.
template <typename Self, typename Batch = std::vector<Example<>>>
using StreamDataset = BatchDataset<Self, Batch, /*BatchRequest=*/size_t>;
} // namespace datasets
} // namespace data
} // namespace torch