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
|
#include <iostream>
#include <gtest/gtest.h>
#include "caffe2/core/context.h"
#include "caffe2/core/context_gpu.h"
#include "caffe2/core/flags.h"
#include "caffe2/operators/reshape_op.h"
#include "caffe2/utils/math.h"
C10_DECLARE_string(caffe_test_root);
namespace caffe2 {
static void AddConstInput(
const vector<int64_t>& shape,
const float value,
const string& name,
Workspace* ws) {
DeviceOption option;
option.set_device_type(PROTO_CUDA);
CUDAContext context(option);
Blob* blob = ws->CreateBlob(name);
auto* tensor = BlobGetMutableTensor(blob, CUDA);
tensor->Resize(shape);
math::Set<float, CUDAContext>(
tensor->numel(), value, tensor->template mutable_data<float>(), &context);
return;
}
TEST(ReshapeOpGPUTest, testReshapeWithScalar) {
if (!HasCudaGPU())
return;
Workspace ws;
OperatorDef def;
def.set_name("test_reshape");
def.set_type("Reshape");
def.add_input("X");
def.add_output("XNew");
def.add_output("OldShape");
def.add_arg()->CopyFrom(MakeArgument("shape", vector<int64_t>{1}));
def.mutable_device_option()->set_device_type(PROTO_CUDA);
AddConstInput(vector<int64_t>(), 3.14, "X", &ws);
// execute the op
unique_ptr<OperatorBase> op(CreateOperator(def, &ws));
EXPECT_TRUE(op->Run());
Blob* XNew = ws.GetBlob("XNew");
const Tensor& XNewTensor = XNew->Get<Tensor>();
EXPECT_EQ(1, XNewTensor.dim());
EXPECT_EQ(1, XNewTensor.numel());
}
} // namespace caffe2
|
