| 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(): |