1 from __future__
import absolute_import
2 from __future__
import division
3 from __future__
import print_function
6 from caffe2.proto
import caffe2_pb2
10 def build_net(net_name, cross_socket):
11 net = core.Net(net_name)
12 net.Proto().type =
"async_scheduling" 13 numa_device_option = caffe2_pb2.DeviceOption()
14 numa_device_option.device_type = caffe2_pb2.CPU
15 numa_device_option.numa_node_id = 0
16 net.XavierFill([], net_name +
"/input_blob", shape=[1024, 1024],
17 device_option=numa_device_option)
19 numa_device_option.numa_node_id = 1
20 net.Copy(net_name +
"/input_blob", net_name +
"/output_blob",
21 device_option=numa_device_option)
26 assert workspace.IsNUMAEnabled()
and workspace.GetNumNUMANodes() >= 2
28 single_net = build_net(
"single_net",
False)
29 cross_net = build_net(
"cross_net",
True)
30 workspace.CreateNet(single_net)
31 workspace.CreateNet(cross_net)
35 workspace.RunNet(single_net.Name(), 5000)
36 print(
"Single socket time:", time.time() - t)
39 workspace.RunNet(cross_net.Name(), 5000)
40 print(
"Cross socket time:", time.time() - t)
43 if __name__ ==
'__main__':
44 core.GlobalInit([
"caffe2",
"--caffe2_cpu_numa_enabled=1"])