Compute the lipschitz of a deep network effectively

KaiHoo: <p>Hi, I want to compute the lipschitz constant of a network, but it seems slow. Here is how I implement it:</p>
<pre><code class="lang-auto"># naive version
lc = 1
for block in self.blocks:
    lc = lc * compute_lipschitz(block)
</code></pre>
<p>The function <code>compute_lipschitz</code> is like:</p>
<pre><code class="lang-auto">def compute_lipschitz(conv)：
    x = torch.randn_like(input_of_conv)
    for i in range(5):
        x = F.conv(x, conv.weight, stride, padding, ...)
        x = F.conv_tranpose(x, conv.weight, stride, padding, ...)
        x = x / x.norm()
    return F.conv(x, ...).norm()
</code></pre>
<p>for convolution layers. For activation layers, it returns the maximum absolute derivative, e.g., for relu, it returns 1.</p>
<p>However, I find it very slow since 1) the network is deep, 2) the computation in <code>compute_lipschitz</code> is not heavy and cannot take full benefits of GPUs, 3) The computation has to be repeated on all GPUs if I use DDP, which I think is not very necessary.</p>
<p>I tried two ideas to solve this problem, but they did not work.</p>
<ol>
<li>
<p>Use two CUDA stream, one for the model forward, one for the <code>compute_lipschitz</code>, it seems not very fast.</p>
</li>
<li>
<p>With DDP on 8 GPUs, one GPU only compute the lipschitz of 1/8 blocks:</p>
</li>
</ol>
<pre><code class="lang-auto"># accelerate version
lc = torch.zeros(len(self.blocks)).to(gpu_device)
for idx, block in enumerate(self.blocks):
    if idx % num_gpus = rank:
        lc[idx] = compute_lipschitz(block)

torch.distributed.all_reduce(lc)
lc = lc.prod()
</code></pre>
<p>This method is 8 times faster and the result is correct as the naive version. However, the gradient seems not correct. The training of my network is very different as the naive version, though the value lc is computed correctly.</p>
<p>Any ideas to solve the gradient problem or accelerate the naive version? Thanks!</p>

agu: <p>This is a very interesting problem <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Some initial clarifications:</p>
<ol>
<li>More precisely, is <code>lc = torch.ones(1)</code> in the naive version? (You mention that it requires gradient, so I assume it is a <code>torch.Tensor</code>.)</li>
<li>Reading the docs for <code>torch.nn.utils.spectral_norm</code>, the function appears to modify the parameter in-place. Does <code>compute_lipschitz</code> do any in-place modification, or does it simply return a value? If it returns a value, does it return a Python <code>float</code>, a <code>torch.Tensor</code>, or something else?</li>
<li>Is there any data dependency between the model forward and computing the model Lipschitz constant?</li>
</ol>
<p>Also, my guess is that the gradient is not correct because the autograd history of <code>lc = torch.zeros(len(self.blocks))</code> is not propagated through the all-reduce. Each rank only has the history for its fraction of the elements, so gradients will not flow the loss back through to <code>lc</code> as expected.</p>
<p>Finally, is there any way to write a batched <code>compute_lipschitz(blocks)</code> that can perform the <code>prod</code> in the kernel directly?</p>

agu: <p>It looks like if you use this <a href="https://github.com/pytorch/pytorch/blob/efccb6401c6451389c9005a43c29fd055fb89452/torch/distributed/nn/functional.py#L207" rel="noopener nofollow ugc"><code>all_reduce</code> API</a>, the autograd history will be propagated appropriately (by calling an additional all-reduce in the backward). Maybe try dropping this in instead of <code>dist.all_reduce</code> and see if it fixes your problem.</p>

KaiHoo: <p>Hi <a class="mention" href="/u/agu">@agu</a> Thanks for your reply. I made some modification to the question, so it can be more clear.</p>
<ol>
<li>
<p>For <code>lc = torch.ones(1)</code> it is not in the naive version. In the naive version, I initialized it with 1, and muliply it with all blocks’ lipschitz.</p>
</li>
<li>
<p>The <code>compute_lipschitz</code> does not do any in-place operations, just return a number with gradient if the computed module has parameters, e.g., conv layer. Or a number without gradient it the  computed module has no parameters, e.g., ReLU layer .</p>
</li>
<li>
<p>The computing of the model Lipschitz is independent of the model forward.</p>
</li>
</ol>
<p>I will see if I can do a batched <code>compute_lipschitz</code> for layers whose parameters’ shape are the same.<br>
I wonder what is the difference between <code>torch.distributed.all_reduce</code> and  <code>torch.distributed.nn.functional.all_reduce</code>? Thanks!</p>

agu: <blockquote>
<p>I wonder what is the difference between <code>torch.distributed.all_reduce</code> and <code>torch.distributed.nn.functional.all_reduce</code> ? Thanks!</p>
</blockquote>
<p>Sorry, if I was unclear. The difference is that the <code>nn.functional.all_reduce</code> one should propagate the gradients from all ranks, while the normal <code>distributed.all_reduce</code> does not. In my understanding, if you have a tensor <code>t = torch.zeros(world_size)</code>, run <code>t[rank] = &lt;some computation tracked by autograd&gt;</code> on each rank, call <code>nn.functional.all_reduce(t)</code> on each rank, and use the resulting <code>t</code> for loss computation, then calling <code>loss.backward()</code> will provide gradients for all <code>world_size</code> elements in <code>t</code>, not only just <code>t[rank]</code>. This should be because the <code>nn.functional.all_reduce</code> will additionally <a href="https://github.com/pytorch/pytorch/blob/efccb6401c6451389c9005a43c29fd055fb89452/torch/distributed/nn/functional.py#L445" rel="noopener nofollow ugc">all-reduce each rank’s gradient</a>.</p>
<p>Regarding your answers to 1 and 2, my concern is that I do not think that gradients can propagate from scalars (i.e. Python <code>float</code>) back through to <code>torch.Tensor</code>s. Either way, I still suggest trying out <code>nn.functional.all_reduce</code> first and see if that works.</p>

KaiHoo: <p><a class="mention" href="/u/agu">@agu</a> After using <code>torch.distributed.nn.functional.all_reduce</code> , the program stuck at the first iteration <img src="https://discuss.pytorch.org/images/emoji/apple/upside_down_face.png?v=12" title=":upside_down_face:" class="emoji" alt=":upside_down_face:" loading="lazy" width="20" height="20"></p>

agu: <p><img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji only-emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<p>Would it be possible to share a minimal version of your script? I can try to take a look.</p>

KaiHoo: <p>Sure I will make it today.</p>

KaiHoo: <p><a class="mention" href="/u/agu">@agu</a> Sorry for the late reply. I was busy yesterday. Here is a minimal version. I removed some complicated loss so that the main idea is clear.</p>
<p>The problem is solved by batching as you suggested. The runtime of my model is reduced by 40%. But it is still interesting to see why all_reduce not works here.</p>
<pre><code class="lang-auto">import argparse
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.nn.functional as F


class conv3x3(nn.Conv2d):
    def __init__(self,
                 in_channels,
                 out_channels,
                 n_power_iterations=10,
                 eps=1e-12,
                 input_size=32) -&gt; None:
        super(conv3x3, self).__init__(in_channels=in_channels,
                                      out_channels=out_channels,
                                      kernel_size=3,
                                      stride=1,
                                      padding=1,
                                      bias=True)
        nn.init.orthogonal_(self.weight)

        self.num = n_power_iterations
        self.eps = eps

        init_x = torch.randn(1, self.in_channels, input_size, input_size)
        self.register_buffer('init_x', init_x)

    def lipschitz(self):
        x = self.init_x.data

        # Power Method to find the lipschitz
        for _ in range(self.num_iter):
            x = F.conv2d(x, self.weight, bias=None, padding=1)
            x = F.conv_transpose2d(x, self.weight, bias=None, padding=1)
            x = x / (x.norm() + self.eps)

        self.init_x += (x - self.init_x).detach()
        x = F.conv2d(x,
                     self.weight,
                     bias=None,
                     stride=self.stride,
                     padding=self.padding)
        return x.norm()


class ReLU(nn.ReLU):
    def __init__(self):
        super(ReLU, self).__init__()

    def lipschitz(self):
        return 1.


class ToyModel(nn.Module):
    def __init__(self, depth):
        layers = []
        for _ in range(depth):
            layers.append(conv3x3(64, 64))
            layers.append(ReLU())

        self.layers = layers

    def forward(self, x):
        out = self.layers(x)
        lc = self.lipschitz()
        return out, lc

    def lipschitz(self):
        if dist.is_initialized():
            return self.fast_lipschitz()
        lc = 1
        for module in self.layers:
            lc = lc * module.lipschitz()
        return lc

    def fast_lipschitz(self):
        lc_list = torch.zeros(len(self.layers))
        lc_list = lc_list.to(self.layers[0].weight.device)

        rank = dist.get_rank()
        num_gpus = dist.get_world_size()

        for idx, module in enumerate(self.layers):
            if idx % num_gpus == rank:
                lc_list[idx] = module.lipschitz()

        dist.all_reduce(lc_list)
        return lc_list.prod()


if __name__ == '__main__':
    dist.init_process_group(backend="nccl")
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int)
    args = parser.parse_args()

    local_rank = args.local_rank
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    model = ToyModel(20).to(device)
    model = nn.parallel.DistributedDataParallel(model)

    inputs = torch.randn(8, 64, 32, 32).to(device)
    targets = torch.randn(8, 64, 32, 32).to(device)

    optimizer = torch.optim.Adam(model.parameters())

    for _ in range(100):
        optimizer.zero_step()
        outputs, lc = model(inputs)
        loss = (outputs - targets).pow(2).mean() + lc * 0.1
        loss.backward()
        optimizer.step()

</code></pre>

KaiHoo: <aside class="quote no-group" data-username="agu" data-post="2" data-topic="163025">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/agu/48/49913_2.png" class="avatar"> agu:</div>
<blockquote>
<p>Finally, is there any way to write a batched <code>compute_lipschitz(blocks)</code> that can perform the <code>prod</code> in the kernel directly?</p>
</blockquote>
</aside>
<p>Today I tried this and it is really fast. The lipschtiz constants  of layers have the same shape of weights can be computed together. Say my residual network has 3 stages, and in each stage the convolutional layers’ shapes are the same. I just concatenate the weights of these layers (as a single group convolution) and do the computation together. Although the computation cost is not changed, the batching idea is much faster.</p>