Caffe2 - Python API
A deep learning, cross platform ML framework
numa_benchmark.py
1 from __future__ import absolute_import
2 from __future__ import division
3 from __future__ import print_function
4 
5 from caffe2.python import core, workspace
6 from caffe2.proto import caffe2_pb2
7 import time
8 
9 SHAPE_LEN = 4096
10 NUM_ITER = 1000
11 GB = 1024 * 1024 * 1024
12 NUM_REPLICAS = 48
13 
14 
15 def build_net(net_name, cross_socket):
16  init_net = core.Net(net_name + "_init")
17  init_net.Proto().type = "async_scheduling"
18  numa_device_option = caffe2_pb2.DeviceOption()
19  numa_device_option.device_type = caffe2_pb2.CPU
20  numa_device_option.numa_node_id = 0
21  for replica_id in range(NUM_REPLICAS):
22  init_net.XavierFill([], net_name + "/input_blob_" + str(replica_id),
23  shape=[SHAPE_LEN, SHAPE_LEN], device_option=numa_device_option)
24 
25  net = core.Net(net_name)
26  net.Proto().type = "async_scheduling"
27  if cross_socket:
28  numa_device_option.numa_node_id = 1
29  for replica_id in range(NUM_REPLICAS):
30  net.Copy(net_name + "/input_blob_" + str(replica_id),
31  net_name + "/output_blob_" + str(replica_id),
32  device_option=numa_device_option)
33  return init_net, net
34 
35 
36 def main():
37  assert workspace.IsNUMAEnabled() and workspace.GetNumNUMANodes() >= 2
38 
39  single_init, single_net = build_net("single_net", False)
40  cross_init, cross_net = build_net("cross_net", True)
41 
42  workspace.CreateNet(single_init)
43  workspace.RunNet(single_init.Name())
44  workspace.CreateNet(cross_init)
45  workspace.RunNet(cross_init.Name())
46 
47  workspace.CreateNet(single_net)
48  workspace.CreateNet(cross_net)
49 
50  for _ in range(4):
51  t = time.time()
52  workspace.RunNet(single_net.Name(), NUM_ITER)
53  dt = time.time() - t
54  print("Single socket time:", dt)
55  single_bw = 4 * SHAPE_LEN * SHAPE_LEN * NUM_REPLICAS * NUM_ITER / dt / GB
56  print("Single socket BW: {} GB/s".format(single_bw))
57 
58  t = time.time()
59  workspace.RunNet(cross_net.Name(), NUM_ITER)
60  dt = time.time() - t
61  print("Cross socket time:", dt)
62  cross_bw = 4 * SHAPE_LEN * SHAPE_LEN * NUM_REPLICAS * NUM_ITER / dt / GB
63  print("Cross socket BW: {} GB/s".format(cross_bw))
64  print("Single BW / Cross BW: {}".format(single_bw / cross_bw))
65 
66 
67 if __name__ == '__main__':
68  core.GlobalInit(["caffe2", "--caffe2_cpu_numa_enabled=1"])
69  main()