reid from https://github.com/michuanhaohao/reid-strong-baseline
zhangmeng
2020-01-11 bdf3ad71583fb4ef100d3819ecdae8fd9f70083e
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
#pragma once
 
#include <torch/data/example.h>
#include <torch/data/transforms/lambda.h>
 
#include <vector>
 
namespace torch {
namespace data {
namespace transforms {
 
/// A `Collation` is a transform that reduces a batch into a single value.
/// The result is a `BatchDataset` that has the type of the single value as its
/// `BatchType`.
template <typename T, typename BatchType = std::vector<T>>
using Collation = BatchTransform<BatchType, T>;
 
/// A `Collate` allows passing a custom function to reduce/collate a batch
/// into a single value. It's effectively the lambda version of `Collation`,
/// which you could subclass and override `operator()` to achieve the same.
///
/// \rst
/// .. code-block:: cpp
///   using namespace torch::data;
///
///   auto dataset = datasets::MNIST("path/to/mnist")
///     .map(transforms::Collate<Example<>>([](std::vector<Example<>> e) {
///       return std::move(e.front());
///     }));
/// \endrst
template <typename T, typename BatchType = std::vector<T>>
using Collate = BatchLambda<BatchType, T>;
} // namespace transforms
} // namespace data
} // namespace torch