MCPcopy Index your code
hub / github.com/pytorch/pytorch / _GPUInterDeviceBatchNormalization

Function _GPUInterDeviceBatchNormalization

caffe2/python/data_parallel_model.py:2072–2221  ·  view source on GitHub ↗
(model)

Source from the content-addressed store, hash-verified

2070
2071
2072def _GPUInterDeviceBatchNormalization(model):
2073 orig_ops = list(model.net.Proto().op)
2074 new_ops = []
2075 num_devices = len(model._devices)
2076 batch_norm_ops = []
2077 injected_ops = []
2078
2079 spatial_bn_phase = False
2080 sums_blobs = []
2081 sumsq_blobs = []
2082 name = []
2083 input_blob_name = None
2084
2085 spatial_bn_gradient_phase = False
2086 scale_grad_blobs = []
2087 bias_grad_blobs = []
2088 master_device = "cpu_0"
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(

Callers 1

ParallelizeFunction · 0.85

Calls 6

listFunction · 0.85
_gpuReduceFunction · 0.85
stripBlobNameFunction · 0.85
ProtoMethod · 0.45
extendMethod · 0.45
appendMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…