How to average over different/multiple axes of weights matrix?

[jvwilliams23]

Hi again,

I am trying to convert the following from pytorch to LBANN:

# w is weights of a 2D convolution layer, shape is [batch, 512, 512, 3, 3]
# x is a featuremap, shape is [batch, 512, 8, 8]
dcoefs = (w.square().sum(dim=[2, 3, 4]) + 1e-8).rsqrt()  # [NO]
# dcoefs is essentially a scaling for each output channel, shape is [batch, 512]
x = x * dcoefs.reshape(batch_size, -1, 1, 1)

But, I am not sure how to do the ".sum(dim=[2,3,4])" operation.

I tried the following:

• lbann.Reduction, but that does not operate on different axes (output would be a single scalar obtained by summing over all output channels).
• lbann.RowwiseWeightsNorms is not suitable, because I have already applied a transformation to the weights so we get the following error message row-wise weights norms layer "layer378" expects 0 parent layers, but found 1 (layer377).

Any idea how to do this?

Regards,
Josh

[szaman19]

We're working on expanding the Reduction layer but in the meantime, this is the hack I have to do reductions on 3/4D inputs using grouped convolutions with frozen constant weights.

weights = lbann.Weights(initializer=lbann.ConstantInitializer(value=1, optimizer=lbann.NoOptimizer()))
reduction_kernel = lbann.ConvolutionModule(
                                  weights=weights,
                                  num_dims=3, 
                                  kernel=[512, 3 , 3],
                                  groups=512,
                                  output_channels=512,
                                  stride=1,
                                  padding=0,
                                  has_bias=False )
# w has shape (512, 512, 3, 3)
w_reduced = reduction_kernel(w)  # shape (512, 1, 1, 1)

[jvwilliams23]

Works great, thanks! Do you also find that it slows training down quite a lot? I split into two separate convs, still same performance drop:

weights_ones_0 = lbann.Weights(initializer=lbann.ConstantInitializer(value=1), optimizer=lbann.NoOptimizer())
weights_ones_1 = lbann.Weights(initializer=lbann.ConstantInitializer(value=1), optimizer=lbann.NoOptimizer())
reduction_kernel_0 = lm.Convolution3dModule(
    weights=weights_ones_0,
    kernel_size=[1, 3, 3],
    groups=1,
    out_channels=out_channels,
    stride=1,
    padding=0,
    bias=False
)
reduction_kernel_1 = lm.Convolution3dModule(
    weights=weights_ones_1,
    kernel_size=[in_channels, 1, 1],
    groups=out_channels,
    out_channels=out_channels,
    stride=1,
    padding=0,
    bias=False
)
dcoefs = lbann.Rsqrt(reduction_kernel_1(reduction_kernel_0(lbann.Square(w))))

[szaman19]

I haven't seen any significant performance issues when doing this. Our convolution wraps around the DNN library (CuDNN or MIOppen) so it should be fairly performant. Could you let me know how much time the reduction adds?

[jvwilliams23]

Seems to just be backprop on the first epoch actually (test with 1 epoch and the 3D convolutions as shown above):

Label                               | Samples  | Total      | Mean       | Min        | Max        |
  forward prop*                     |        1 | 2.9540e+00 | 2.9540e+00 | 2.9540e+00 | 2.9540e+00 |
  back prop*                        |        1 | 3.6472e+01 | 3.6472e+01 | 3.6472e+01 | 3.6472e+01 |

compared to without:

  forward prop*                     |        1 | 2.7907e+00 | 2.7907e+00 | 2.7907e+00 | 2.7907e+00 |
  back prop*                        |        1 | 2.0235e+00 | 2.0235e+00 | 2.0235e+00 | 2.0235e+00 |

After that, it returns to normal.
100 epochs with 3D conv:

  forward prop*                     |      100 | 1.8650e+01 | 1.8650e-01 | 9.1960e-02 | 2.9164e+00 |
  back prop*                        |      100 | 4.3653e+01 | 4.3653e-01 | 7.1836e-02 | 3.6459e+01 |

100 epochs without:

  forward prop*                     |      100 | 1.7991e+01 | 1.7991e-01 | 1.5099e-01 | 2.7654e+00 |
  back prop*                        |      100 | 8.4930e+00 | 8.4930e-02 | 6.4171e-02 | 2.0390e+00 |

[jvwilliams23]

Hi @szaman19 , how would I modify your hack to sum over the weights of the output channel, rather than the input channel? Something like this:

weights = lbann.Weights(initializer=lbann.ConstantInitializer(value=1, optimizer=lbann.NoOptimizer()))
reduction_kernel = ??
# w has shape (512, 512, 3, 3)
w_reduced = reduction_kernel(w)  # shape (1, 512, 1, 1), NOT (512, 1, 1, 1) as before