pytorch

1

2

3

4

5
from caffe2.python import core, workspace
6
from caffe2.proto import caffe2_pb2
7
from caffe2.python.test_util import TestCase
8
import unittest
9

10
core.GlobalInit(["caffe2", "--caffe2_cpu_numa_enabled=1"])
11

12
def build_test_net(net_name):
13
    net = core.Net(net_name)
14
    net.Proto().type = "async_scheduling"
15

16
    numa_device_option = caffe2_pb2.DeviceOption()
17
    numa_device_option.device_type = caffe2_pb2.CPU
18
    numa_device_option.numa_node_id = 0
19

20
    net.ConstantFill([], "output_blob_0", shape=[1], value=3.14,
21
                         device_option=numa_device_option)
22

23
    numa_device_option.numa_node_id = 1
24
    net.ConstantFill([], "output_blob_1", shape=[1], value=3.14,
25
                         device_option=numa_device_option)
26

27
    gpu_device_option = caffe2_pb2.DeviceOption()
28
    gpu_device_option.device_type = caffe2_pb2.CUDA
29
    gpu_device_option.device_id = 0
30

31
    net.CopyCPUToGPU("output_blob_0", "output_blob_0_gpu",
32
                        device_option=gpu_device_option)
33
    net.CopyCPUToGPU("output_blob_1", "output_blob_1_gpu",
34
                        device_option=gpu_device_option)
35

36
    return net
37

38

39
@unittest.skipIf(not workspace.IsNUMAEnabled(), "NUMA is not enabled")
40
@unittest.skipIf(workspace.GetNumNUMANodes() < 2, "Not enough NUMA nodes")
41
@unittest.skipIf(not workspace.has_gpu_support, "No GPU support")
42
class NUMATest(TestCase):
43
    def test_numa(self):
44
        net = build_test_net("test_numa")
45

46
        workspace.RunNetOnce(net)
47

48
        self.assertEqual(workspace.GetBlobNUMANode("output_blob_0"), 0)
49
        self.assertEqual(workspace.GetBlobNUMANode("output_blob_1"), 1)
50

51

52
if __name__ == '__main__':
53
    unittest.main()
54