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
|
from hypothesis import given
import numpy as np
import hypothesis.strategies as st
from caffe2.python import core, workspace
import caffe2.python.hypothesis_test_util as hu
@st.composite
def _dev_options(draw):
op_dev = draw(st.sampled_from(hu.device_options))
if op_dev == hu.cpu_do:
# the CPU op can only handle CPU tensor
input_blob_dev = hu.cpu_do
else:
input_blob_dev = draw(st.sampled_from(hu.device_options))
return op_dev, input_blob_dev
class TestEnsureCPUOutputOp(hu.HypothesisTestCase):
@given(
input=hu.tensor(dtype=np.float32),
dev_options=_dev_options()
)
def test_ensure_cpu_output(self, input, dev_options):
op_dev, input_blob_dev = dev_options
net = core.Net('test_net')
data = net.GivenTensorFill(
[],
["data"],
values=input,
shape=input.shape,
device_option=input_blob_dev
)
data_cpu = net.EnsureCPUOutput(
[data],
["data_cpu"],
device_option=op_dev
)
workspace.RunNetOnce(net)
data_cpu_value = workspace.FetchBlob(data_cpu)
np.testing.assert_allclose(input, data_cpu_value)
|
