Caffe2 - C++ API A deep learning, cross platform ML framework
numpy_tile_op.h
1 #ifndef CAFFE2_OPERATORS_NUMPY_TILE_OP_H_
2 #define CAFFE2_OPERATORS_NUMPY_TILE_OP_H_
3
4 #include "caffe2/core/common_omp.h"
5 #include "caffe2/core/context.h"
6 #include "caffe2/core/logging.h"
7 #include "caffe2/core/operator.h"
8 #include "caffe2/utils/math.h"
9
10 namespace caffe2 {
11
12 // Copy a Blob n times along a specified axis.
13 template <class Context>
14 class NumpyTileOp : public Operator<Context> {
15  public:
16  USE_OPERATOR_CONTEXT_FUNCTIONS;
17  template <class... Args>
18  explicit NumpyTileOp(Args&&... args)
19  : Operator<Context>(std::forward<Args>(args)...) {}
20  ~NumpyTileOp() {}
21
22  bool RunOnDevice() override {
23  const auto& input = Input(0);
24  const auto& repeats = Input(1);
25
26  // Check that the `repeats` tensor has the correct rank, has a number of
27  // elements equal to the number of axes of `input`.
28  CAFFE_ENFORCE_EQ(repeats.dim(), 1, "repeats input must be a 1-d tensor");
29  CAFFE_ENFORCE_EQ(
30  repeats.numel(),
31  input.dim(),
32  "repeats input have the same"
33  " number of elements as `inputs` has dimensions.");
34  const int64_t *repeats_data = repeats.template data<int64_t>();
35  for (size_t i = 0; i < repeats.numel(); ++i) {
36  CAFFE_ENFORCE_GE(repeats_data[i], 0);
37  }
38
39  auto* output = Output(0);
40
41  // Alternate inputs and outputs between two buffers. Repeatedly apply the
42  // Tile kernel along each axis. Then copy out the resulting data into the
43  // output tensor.
44  Tensor *src = &buffer, *dst = output;
45  src->CopyFrom(input);
46  vector<int64_t> output_dims(input.sizes().vec());
47  for (size_t i = 0; i < repeats.numel(); ++i) {
48  if (repeats_data[i] == 1) {
49  continue;
50  }
51  // size up to (and not including) axis
52  const auto outer_dim = src->size_to_dim(i);
53  // size from axis up
54  const auto inner_dim = src->size_from_dim(i);
55
56  dst->Resize(outer_dim, inner_dim * repeats_data[i]);
57
66  const char* src_data = static_cast<const char*>(src->raw_data());
67  char* dst_data = static_cast<char*>(dst->raw_mutable_data(src->dtype()));
68
69  DoTile(
70  src->dtype(),
71  src->itemsize(),
72  outer_dim,
73  inner_dim,
74  repeats_data[i],
75  src_data,
76  dst_data);
77
78  output_dims[i] *= repeats_data[i];
79  dst->Reshape(output_dims);
80
81  std::swap(src, dst);
82  }
83
84  // NB: because we have the swap at the end of the above loop, our real
85  // result tensor is going to live in *src when we reach this line
86  // whether we entered the loop or not :)
87  if (output != src)
88  output->CopyFrom(*src);
89
90  return true;
91  }
92
93  private:
94  void DoTile(
95  const TypeMeta& meta,
96  int item_size,
97  int outer_dim,
98  int inner_dim,
99  int64_t num_tiles,
100  const char* input_data,
101  char* output_data) {
102  for (auto i = 0; i < outer_dim; ++i) {
103  for (auto t = 0; t < num_tiles; ++t) {
104  context_.CopyItemsSameDevice(meta, inner_dim, input_data, output_data);
105  output_data += inner_dim * item_size;
106  }
107  input_data += inner_dim * item_size;
108  }
109  }
110
111  Tensor buffer{Context::GetDeviceType()};
112 };
113
114 } // namespace caffe2
115
116 #endif // CAFFE2_OPERATORS_NUMPY_TILE_OP_H_
const Tensor & Input(int idx, DeviceType type=Context::GetDeviceType())
Retrieve a non-owning reference to the input at position &#39;idx&#39; for this operator. ...
Definition: operator.h:702
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
TypeMeta is a thin class that allows us to store the type of a container such as a blob...
Definition: typeid.h:324
bool RunOnDevice() override
Definition: numpy_tile_op.h:22