Reduces results from multiple gpus and distributes the results back to each device. This is done by copying values to the master device and summing them. The master device result is then copied back to each of the devices. param: the name of the data (blobs)
(param, num_devices, master_device, result_blobs=None)
| 2089 | master_device_option = core.DeviceOption(caffe2_pb2.CPU) |
| 2090 | |
| 2091 | def _gpuReduce(param, num_devices, master_device, result_blobs=None): |
| 2092 | """ |
| 2093 | Reduces results from multiple gpus and distributes the results back |
| 2094 | to each device. This is done by copying values to the master device |
| 2095 | and summing them. The master device result is then copied back to |
| 2096 | each of the devices. |
| 2097 | |
| 2098 | param: the name of the data (blobs) to reduce |
| 2099 | num_devices: the number of devices |
| 2100 | master_device: the device to copy/compute values on |
| 2101 | result_blobs: optional list of result blobs to copy to |
| 2102 | """ |
| 2103 | added_ops = [] |
| 2104 | source_blobs = [] |
| 2105 | destination_blobs = [] |
| 2106 | if result_blobs is None: |
| 2107 | result_blobs = [ |
| 2108 | "gpu_{}/{}_combined".format(i, param) for i in range(num_devices) |
| 2109 | ] |
| 2110 | for i in range(num_devices): |
| 2111 | device_option = core.DeviceOption(model._device_type, i) |
| 2112 | source_blobs.append("gpu_{}/{}".format(i, param)) |
| 2113 | destination_blobs.append( |
| 2114 | "{}/{}_gpu_{}_copy".format(master_device, param, i)) |
| 2115 | added_ops.append( |
| 2116 | core.CreateOperator( |
| 2117 | "CopyGPUToCPU", |
| 2118 | source_blobs[i], |
| 2119 | destination_blobs[i], |
| 2120 | device_option=device_option)) |
| 2121 | added_ops.append( |
| 2122 | core.CreateOperator( |
| 2123 | "Sum", |
| 2124 | destination_blobs, |
| 2125 | "{}/{}_combined".format(master_device, param), |
| 2126 | device_option=master_device_option)) |
| 2127 | for i in range(num_devices): |
| 2128 | device_option = core.DeviceOption(model._device_type, i) |
| 2129 | added_ops.append( |
| 2130 | core.CreateOperator( |
| 2131 | "CopyCPUToGPU", |
| 2132 | "{}/{}_combined".format(master_device, param), |
| 2133 | result_blobs[i], |
| 2134 | device_option=device_option)) |
| 2135 | return added_ops |
| 2136 | |
| 2137 | for op in orig_ops: |
| 2138 | if op.type != 'SpatialBN' and op.type != 'SpatialBNGradient': |
no test coverage detected
searching dependent graphs…