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

Function _CPUInterDeviceBatchNormalization

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

Source from the content-addressed store, hash-verified

1965
1966
1967def _CPUInterDeviceBatchNormalization(model):
1968 orig_ops = list(model.net.Proto().op)
1969 new_ops = []
1970 num_devices = len(model._devices)
1971 batch_norm_ops = []
1972 injected_ops = []
1973
1974 spatial_bn_phase = False
1975 sums_blobs = []
1976 sumsq_blobs = []
1977 name = []
1978 input_blob_name = None
1979
1980 spatial_bn_gradient_phase = False
1981 scale_grad_blobs = []
1982 bias_grad_blobs = []
1983
1984 def _cpuReduce(param, input_blobs, destination_blobs):
1985 """
1986 Reduce results from multiple cpus and distributes the results back
1987 to each device. This is done by copying values to cpu_0 and summing
1988 them. The cpu_0 result is then copied back to each of the devices.
1989
1990 param: the name of the data (blobs) to reduce
1991 input_blobs: the list of blobs to reduce
1992 destination_blobs: list of blobs to copy the result to
1993 """
1994 added_ops = []
1995 result_blob = "cpu_0/" + param + "_combined"
1996 added_ops.append(core.CreateOperator("Sum", input_blobs, result_blob))
1997 for blob in destination_blobs:
1998 added_ops.append(core.CreateOperator("Copy", result_blob, blob))
1999 return added_ops
2000
2001 for op in orig_ops:
2002 if op.type != 'SpatialBN' and op.type != 'SpatialBNGradient':
2003 if spatial_bn_phase:
2004 new_ops.extend(injected_ops)
2005 new_ops.append(
2006 core.CreateOperator("Sum",
2007 sums_blobs,
2008 input_blob_name + "_sums_combined"))
2009 new_ops.append(
2010 core.CreateOperator("Sum",
2011 sumsq_blobs,
2012 input_blob_name + "_sumsq_combined"))
2013 new_ops.extend(batch_norm_ops)
2014 injected_ops = []
2015 batch_norm_ops = []
2016 sums_blobs = []
2017 sumsq_blobs = []
2018 spatial_bn_phase = False
2019 input_blob_name = None
2020 elif spatial_bn_gradient_phase:
2021 new_ops.extend(injected_ops)
2022 new_ops.extend(_cpuReduce(
2023 stripBlobName(scale_grad_blobs[0]),
2024 scale_grad_blobs,

Callers 1

ParallelizeFunction · 0.85

Calls 6

listFunction · 0.85
_cpuReduceFunction · 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…