mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
9e71de398b
Adding NUMA awareness through numa_node_id in DeviceOption. Blobs of operators with numa_node_id are allocated on corr. memory banks, using CPU pools with NUMA affinity set to run operators.
46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
from caffe2.python import core, workspace
|
|
from caffe2.proto import caffe2_pb2
|
|
import time
|
|
|
|
|
|
def build_net(net_name, cross_socket):
|
|
net = core.Net(net_name)
|
|
net.Proto().type = "async_scheduling"
|
|
numa_device_option = caffe2_pb2.DeviceOption()
|
|
numa_device_option.device_type = caffe2_pb2.CPU
|
|
numa_device_option.numa_node_id = 0
|
|
net.XavierFill([], net_name + "/input_blob", shape=[1024, 1024],
|
|
device_option=numa_device_option)
|
|
if cross_socket:
|
|
numa_device_option.numa_node_id = 1
|
|
net.Copy(net_name + "/input_blob", net_name + "/output_blob",
|
|
device_option=numa_device_option)
|
|
return net
|
|
|
|
|
|
def main():
|
|
assert workspace.IsNUMAEnabled() and workspace.GetNumNUMANodes() >= 2
|
|
|
|
single_net = build_net("single_net", False)
|
|
cross_net = build_net("cross_net", True)
|
|
workspace.CreateNet(single_net)
|
|
workspace.CreateNet(cross_net)
|
|
|
|
for _ in range(4):
|
|
t = time.time()
|
|
workspace.RunNet(single_net.Name(), 5000)
|
|
print("Single socket time:", time.time() - t)
|
|
|
|
t = time.time()
|
|
workspace.RunNet(cross_net.Name(), 5000)
|
|
print("Cross socket time:", time.time() - t)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
core.GlobalInit(["caffe2", "--caffe2_cpu_numa_enabled=1"])
|
|
main()
|