Build a graph containing a sequence of batch normalizations. Args: device: string, the device to run on. input_shape: shape of the input tensor. axes: axes that are to be normalized across. num_layers: number of batch normalization layers in the graph. mode: "op", "py" or "slo
(device, input_shape, axes, num_layers, mode, scale, train)
| 66 | |
| 67 | |
| 68 | def build_graph(device, input_shape, axes, num_layers, mode, scale, train): |
| 69 | """Build a graph containing a sequence of batch normalizations. |
| 70 | |
| 71 | Args: |
| 72 | device: string, the device to run on. |
| 73 | input_shape: shape of the input tensor. |
| 74 | axes: axes that are to be normalized across. |
| 75 | num_layers: number of batch normalization layers in the graph. |
| 76 | mode: "op", "py" or "slow" depending on the implementation. |
| 77 | scale: scale after normalization. |
| 78 | train: if true, also run backprop. |
| 79 | |
| 80 | Returns: |
| 81 | An array of tensors to run() |
| 82 | """ |
| 83 | moment_shape = [] |
| 84 | keep_dims = mode == "py" or mode == "slow" |
| 85 | if keep_dims: |
| 86 | for axis in range(len(input_shape)): |
| 87 | if axis in axes: |
| 88 | moment_shape.append(1) |
| 89 | else: |
| 90 | moment_shape.append(input_shape[axis]) |
| 91 | else: |
| 92 | for axis in range(len(input_shape)): |
| 93 | if axis not in axes: |
| 94 | moment_shape.append(input_shape[axis]) |
| 95 | with ops.device("/%s:0" % device): |
| 96 | tensor = variables.Variable(random_ops.truncated_normal(input_shape)) |
| 97 | for _ in range(num_layers): |
| 98 | if train: |
| 99 | mean, variance = nn_impl.moments(tensor, axes, keep_dims=keep_dims) |
| 100 | else: |
| 101 | mean = array_ops.zeros(moment_shape) |
| 102 | variance = array_ops.ones(moment_shape) |
| 103 | beta = variables.Variable(array_ops.zeros(moment_shape)) |
| 104 | gamma = variables.Variable(constant_op.constant(1.0, shape=moment_shape)) |
| 105 | if mode == "py": |
| 106 | tensor = batch_norm_py(tensor, mean, variance, beta, gamma, scale) |
| 107 | elif mode == "op": |
| 108 | tensor = batch_norm_op(tensor, mean, variance, beta, gamma, scale) |
| 109 | elif mode == "slow": |
| 110 | tensor = batch_norm_slow(tensor, mean, variance, beta, gamma, scale) |
| 111 | if train: |
| 112 | return gradients_impl.gradients([tensor], variables.trainable_variables()) |
| 113 | else: |
| 114 | return [tensor] |
| 115 | |
| 116 | |
| 117 | def print_difference(mode, t1, t2): |
no test coverage detected