100x more time cost of DCGAN Hessian vector product

Hongkai_Zheng: <p>I found a weird computation cost  when I was trying to compute Hessian vector product of DCGAN.<br>
The Hessian vector product for DCGAN costs <strong>100x</strong> more than the other GAN , even though it contains less parameters.  I’m confused because the problem seems not caused by one single module since the other GAN contains linear, pooling, convolution while DCGAN only has convolution. Could anyone figure out what’s going on here?</p>
<p>I use CUDA events to measure time cost of backward of each model and Hessian vector product. And here is the result.</p>
<pre><code class="lang-auto">=====Test1=====
Discriminator backward :  10.1550083ms
Generator backward :  6.1214719ms
Hessian vector product :  13.5014400ms
=====Test2=====
Discriminator backward :  18.8948479ms
Generator backward :  31.2074242ms
Hessian vector product :  1206.1020508ms
</code></pre>
<p>The batchnorms and activation functions don’t make a difference so I removed all of them to make the code clear. The math form of hessian vector product here is<br>
$\frac{\partial^2f}{\partial G\partial D} \frac{\partial f}{\partial D}$</p>
<pre><code class="lang-auto">class DC_generator(nn.Module):
    def __init__(self, z_dim=100, channel_num=3, feature_num=64):
        super(DC_generator, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(z_dim, feature_num * 8, kernel_size=4, stride=1, padding=0, bias=True),
            # (feature_num * 8) * 4x4
            nn.ConvTranspose2d(feature_num * 8, feature_num * 4, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 4) * 8x8
            nn.ConvTranspose2d(feature_num * 4, feature_num * 2, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 2) * 16x16
            nn.ConvTranspose2d(feature_num * 2, feature_num, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 2) * 32x32
            nn.ConvTranspose2d(feature_num, channel_num, kernel_size=4, stride=2, padding=1, bias=True),
            # channel_num * 64x64
        )

    def forward(self, input):
        return self.main(input)


class DC_discriminator(nn.Module):
    def __init__(self, channel_num=3, feature_num=64):
        super(DC_discriminator, self).__init__()
        self.main = nn.Sequential(
            # channel_num * 64x64
            nn.Conv2d(channel_num, feature_num, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num) * 32x32
            nn.Conv2d(feature_num, feature_num * 2, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 2) * 16x16
            nn.Conv2d(feature_num * 2, feature_num * 4, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 4) * 8x8
            nn.Conv2d(feature_num * 4, feature_num * 8, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 8) * 4x4
            nn.Conv2d(feature_num * 8, 1, kernel_size=4, stride=1, padding=0, bias=True),
            # feature_num * 16x16
        )

    def forward(self, input):
        return self.main(input)


class dc_D(nn.Module):
    def __init__(self):
        super(dc_D, self).__init__()
        self.conv = nn.Sequential(
            # 3 * 32x32
            nn.Conv2d(in_channels=3, out_channels=128, kernel_size=5, stride=1),
            nn.LeakyReLU(0.01),
            nn.MaxPool2d(2,2),
            # 32 * 14x14
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=5, stride=1),
            nn.LeakyReLU(0.01),
            nn.MaxPool2d(2,2)
            # 64 * 5x5
        )
        self.fc = nn.Sequential(
            nn.Linear(1600 * 2, 1024),
            nn.LeakyReLU(0.01),
            nn.Linear(1024, 1)
        )

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.shape[0], -1)
        return self.fc(x)


class dc_G(nn.Module):
    def __init__(self, z_dim=96):
        super(dc_G, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(z_dim, 1024),
            nn.Linear(1024, 8 * 8 * 128),
        )
        self.convt = nn.Sequential(
            nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=4, stride=2, padding=1),
            nn.ConvTranspose2d(in_channels=64, out_channels=3, kernel_size=4, stride=2, padding=1)
        )

    def forward(self, x):
        x = self.fc(x)
        x = x.view(x.shape[0], 128, 8, 8)
        return self.convt(x)

def hvptest(D, G, z):
    criterion = nn.BCEWithLogitsLoss()
    d_fake = D(G(z))
    loss = criterion(d_fake, torch.zeros(d_fake.shape, device=device))

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    grad_d = torch.autograd.grad(loss, D.parameters(), create_graph=True, retain_graph=True)
    end.record()
    torch.cuda.synchronize(device=device)
    print('Discriminator backward :  %.7fms' % (start.elapsed_time(end)))
    grad_d_vec = torch.cat([g.contiguous().view(-1) for g in grad_d])

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    grad_g = torch.autograd.grad(loss, G.parameters(), create_graph=True, retain_graph=True)
    end.record()
    torch.cuda.synchronize(device=device)
    print('Generator backward :  %.7fms' % (start.elapsed_time(end)))

    grad_g_vec = torch.cat([g.contiguous().view(-1) for g in grad_g])
    # print('Discriminator parameter number: %d' % grad_d_vec.numel())
    # print('Generator parameter number: %d' % grad_g_vec.numel())
    vec_d = grad_d_vec.clone().detach()

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    hvp_d = torch.autograd.grad(grad_d_vec, G.parameters(), grad_outputs=vec_d, retain_graph=True)
    end.record()
    torch.cuda.synchronize(device=device)
    print('Hessian vector product :  %.7fms' % (start.elapsed_time(end)))


if __name__ == '__main__':
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    batch_size = 128
    print('=====Test1=====')
    D = dc_D().to(device)
    G = dc_G(z_dim=96).to(device)
    img_cifar = (torch.rand((batch_size, 3, 32, 32), device=device) - 0.5) / 0.5
    z_mnist = torch.randn((batch_size, 96), device=device)
    hvptest(D=D, G=G, z=z_mnist)
    
    print('=====Test2=====')
    D = DC_discriminator(channel_num=3, feature_num=64).to(device)
    G = DC_generator(z_dim=100, channel_num=3, feature_num=64).to(device)
    z_celeba = torch.randn((batch_size, 100, 1, 1), device=device)
    hvptest(D=D, G=G, z=z_celeba)
</code></pre>
<p>PS: Pytorch version: 1.1.0, GPU: Tesla V100</p>

albanD: <p>Hi,<br>
Thanks for the repro code !</p>
<p>I ran this on cpu (I don’t have a gpu install at hand) with the following modifications:</p>
<pre><code class="lang-python">import time
def hvptest(D, G, z):
    criterion = nn.BCEWithLogitsLoss()
    d_fake = D(G(z))
    loss = criterion(d_fake, torch.zeros(d_fake.shape, device=device))

    start = time.time()
    grad_d = torch.autograd.grad(loss, D.parameters(), create_graph=True, retain_graph=True)
    end = time.time()
    print('Discriminator backward :  %.7fms' % (end - start))
    grad_d_vec = torch.cat([g.contiguous().view(-1) for g in grad_d])

    start = time.time()
    grad_g = torch.autograd.grad(loss, G.parameters(), create_graph=True, retain_graph=True)
    end = time.time()
    print('Generator backward :  %.7fms' % (end - start))

    grad_g_vec = torch.cat([g.contiguous().view(-1) for g in grad_g])
    # print('Discriminator parameter number: %d' % grad_d_vec.numel())
    # print('Generator parameter number: %d' % grad_g_vec.numel())
    vec_d = grad_d_vec.clone().detach()

    start = time.time()
    hvp_d = torch.autograd.grad(grad_d_vec, G.parameters(), grad_outputs=vec_d, retain_graph=True)
    end = time.time()
    print('Hessian vector product :  %.7fms' % (end - start))


if __name__ == '__main__':
    device = torch.device('cpu')
    batch_size = 128
    print('=====Test1=====')
    D = dc_D().to(device)
    G = dc_G(z_dim=96).to(device)
    img_cifar = (torch.rand((batch_size, 3, 32, 32), device=device) - 0.5) / 0.5
    z_mnist = torch.randn((batch_size, 96), device=device)
    hvptest(D=D, G=G, z=z_mnist)

    print('=====Test2=====')
    D = DC_discriminator(channel_num=3, feature_num=64).to(device)
    G = DC_generator(z_dim=100, channel_num=3, feature_num=64).to(device)
    z_celeba = torch.randn((batch_size, 100, 1, 1), device=device)
    hvptest(D=D, G=G, z=z_celeba)
</code></pre>
<p>And I get`</p>
<pre><code class="lang-auto">=====Test1=====
Discriminator backward :  0.2947176ms
Generator backward :  0.4657667ms
Hessian vector product :  1.2660155ms
=====Test2=====
Discriminator backward :  0.1807017ms
Generator backward :  0.9121838ms
Hessian vector product :  2.6689870ms
</code></pre>
<p>Do you see similar results on cpu?<br>
If you still see the problem, which version of pytorch are you using?</p>

Hongkai_Zheng: <p>Thank you for reply.  I got the similar results on CPU. But I don’t get it. Could you explain a bit more? I still don’t understand why DCGAN has much more time cost of Hessian vector product. (BTW, my Pytorch version is 1.1.0)</p>

albanD: <p>Hi,</p>
<p>Your results are similar to mine? Or similar to your GPU results? <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

Hongkai_Zheng: <p>I tried on CPU too. The result on CPU is similar to yours. But for GPU, we need CUDA event to do accurate time measurement.</p>

albanD: <p>You can use</p>
<pre><code class="lang-auto">torch.cuda.synchronize()
start = time.time()
</code></pre>
<p>everywhere there is a <code>time.time()</code> call to avoid using events. Can you check what is the behavior of this when you run on gpu? Does it give the same result as cpu or the same result as with events?</p>

Hongkai_Zheng: <p>I got the same result as with events.</p>
<pre><code class="lang-auto">=====Test1=====
Discriminator backward :  0.0120358 s
Discriminator backward :  0.0067697 s
Hessian vector product :  0.0130763 s
=====Test2=====
Discriminator backward :  0.0188158 s
Discriminator backward :  0.0291374 s
Hessian vector product :  1.1108479 s
</code></pre>

albanD: <p>Interesting, maybe <a class="mention" href="/u/ngimel">@ngimel</a> has some insight in this? Otherwise I’ll set up a gpu install to test this.</p>

Hongkai_Zheng: <p>Hi, I messaged her but she didn’t reply. could you help me ask <a class="mention" href="/u/ngimel">@ngimel</a> for help?</p>

ngimel: <p>Sorry, did not get notification, and sorry, don’t have insight on this. <a class="mention" href="/u/alband">@albanD</a> can you take a look at this or should I?</p>

ptrblck: <p>In case that helps: Repro on TitanV with Pytorch <code>1.3.0.dev20190910</code> (nightly) with CUDA10.0 and NVIDIA driver 418.56:</p>
<pre><code class="lang-auto">=====Test1=====
Discriminator backward :  0.0056169ms
Generator backward :  0.0103526ms
Hessian vector product :  0.0136104ms
=====Test2=====
Discriminator backward :  0.0149076ms
Generator backward :  0.0242238ms
Hessian vector product :  1.0203860ms
</code></pre>
<p>(just added <code>torch.cuda.synchronize()</code> before starting and stopping the timers, no warmup etc.)</p>

SimonW: <p>I suspect that double backward for conv T is really slow.</p>

ngimel: <p><a class="mention" href="/u/simonw">@simonw</a> but there should be double backward for ConvT both in test1 and test2, and only test2 is painfully slow.<br>
<a class="mention" href="/u/ptrblck">@ptrblck</a> can you get nvvp for it and may be run it through pyprof, though I suspect pyprof will choke on double backward.</p>

ptrblck: <p>Sure.<br>
Rerunning on a V100 32GB (Titan V is currently busy):<br>
Reference (using <a class="mention" href="/u/alband">@albanD</a>’s code + adding <code>torch.cuda.synchronize()</code>):</p>
<pre><code class="lang-auto"> =====Test1=====
Discriminator backward :  0.0079403s
Generator backward :  0.0082126s
Hessian vector product :  0.0118001s
=====Test2=====
Discriminator backward :  0.0121739s
Generator backward :  0.0269434s
Hessian vector product :  0.7245934s
</code></pre>
<p>I’ve changed the output to seconds, as I believe <code>time.time()</code> calculates the seconds not ms.<br>
Attaching <a href="https://gist.github.com/ptrblck/509bdb807262d9186c55396eafc3c691" rel="nofollow noopener"><code>nvprof</code></a>.<br>
Let me know, if you need more profiling output etc.</p>

SimonW: <p>You are right. I just saw that <code>DC_generator</code> has considerably more convTs and was wondering.</p>

ngimel: <p>The solution is simple, use <code>torch.backends.cudnn.benchmark=True</code>. cudnn heuristics picks truly atrocious algorithm for couple layers (one of the kernels takes 600 ms, another 50 ms). When cudnn is actually forced to benchmark and pick the best algo, the hessian computation speed becomes reasonable.</p>

Hongkai_Zheng: <p>Hi, I set benchmark=True but got worse results. Could you share your code?</p>

ngimel: <pre><code class="lang-auto">import torch
import time
import torch.nn as nn
torch.backends.cudnn.benchmark=True
class DC_generator(nn.Module):
    def __init__(self, z_dim=100, channel_num=3, feature_num=64):
        super(DC_generator, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(z_dim, feature_num * 8, kernel_size=4, stride=1, padding=0, bias=True),
            # (feature_num * 8) * 4x4
            nn.ConvTranspose2d(feature_num * 8, feature_num * 4, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 4) * 8x8
            nn.ConvTranspose2d(feature_num * 4, feature_num * 2, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 2) * 16x16
            nn.ConvTranspose2d(feature_num * 2, feature_num, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 2) * 32x32
            nn.ConvTranspose2d(feature_num, channel_num, kernel_size=4, stride=2, padding=1, bias=True),
            # channel_num * 64x64
        )

    def forward(self, input):
        return self.main(input)


class DC_discriminator(nn.Module):
    def __init__(self, channel_num=3, feature_num=64):
        super(DC_discriminator, self).__init__()
        self.main = nn.Sequential(
            # channel_num * 64x64
            nn.Conv2d(channel_num, feature_num, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num) * 32x32
            nn.Conv2d(feature_num, feature_num * 2, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 2) * 16x16
            nn.Conv2d(feature_num * 2, feature_num * 4, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 4) * 8x8
            nn.Conv2d(feature_num * 4, feature_num * 8, kernel_size=4, stride=2, padding=1, bias=True),
            # (feature_num * 8) * 4x4
            nn.Conv2d(feature_num * 8, 1, kernel_size=4, stride=1, padding=0, bias=True),
            # feature_num * 16x16
        )

    def forward(self, input):
        return self.main(input)

class dc_D(nn.Module):
    def __init__(self):
        super(dc_D, self).__init__()
        self.conv = nn.Sequential(
            # 3 * 32x32
            nn.Conv2d(in_channels=3, out_channels=128, kernel_size=5, stride=1),
            nn.LeakyReLU(0.01),
            nn.MaxPool2d(2,2),
            # 32 * 14x14
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=5, stride=1),
            nn.LeakyReLU(0.01),
            nn.MaxPool2d(2,2)
            # 64 * 5x5
        )
        self.fc = nn.Sequential(
            nn.Linear(1600 * 2, 1024),
            nn.LeakyReLU(0.01),
            nn.Linear(1024, 1)
        )

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.shape[0], -1)
        return self.fc(x)


class dc_G(nn.Module):
    def __init__(self, z_dim=96):
        super(dc_G, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(z_dim, 1024),
            nn.Linear(1024, 8 * 8 * 128),
        )
        self.convt = nn.Sequential(
            nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=4, stride=2, padding=1),
            nn.ConvTranspose2d(in_channels=64, out_channels=3, kernel_size=4, stride=2, padding=1)
        )

    def forward(self, x):
        x = self.fc(x)
        x = x.view(x.shape[0], 128, 8, 8)
        return self.convt(x)



def hvptest(D, G, z):
    criterion = nn.BCEWithLogitsLoss()
    d_fake = D(G(z))
    loss = criterion(d_fake, torch.zeros(d_fake.shape, device=device))

    torch.cuda.synchronize()
    start = time.time()
    grad_d = torch.autograd.grad(loss, D.parameters(), create_graph=True, retain_graph=True)
    torch.cuda.synchronize()
    end = time.time()
    print('Discriminator backward :  %.7fms' % (end - start))
    grad_d_vec = torch.cat([g.contiguous().view(-1) for g in grad_d])

    torch.cuda.synchronize()
    start = time.time()
    grad_g = torch.autograd.grad(loss, G.parameters(), create_graph=True, retain_graph=True)
    torch.cuda.synchronize()
    end = time.time()
    print('Generator backward :  %.7fms' % (end - start))

    grad_g_vec = torch.cat([g.contiguous().view(-1) for g in grad_g])
    # print('Discriminator parameter number: %d' % grad_d_vec.numel())
    # print('Generator parameter number: %d' % grad_g_vec.numel())
    vec_d = grad_d_vec.clone().detach()
    torch.cuda.synchronize()
    start = time.time()
    hvp_d = torch.autograd.grad(grad_d_vec, G.parameters(), grad_outputs=vec_d, retain_graph=True)
    torch.cuda.synchronize()
    end = time.time()
    print('Hessian vector product :  %.7fms' % (end - start))


if __name__ == '__main__':
    device = torch.device('cuda')
    batch_size = 128
    print('=====Test1=====')
    D = dc_D().to(device)
    G = dc_G(z_dim=96).to(device)
    img_cifar = (torch.rand((batch_size, 3, 32, 32), device=device) - 0.5) / 0.5
    z_mnist = torch.randn((batch_size, 96), device=device)
    for i in range(10):
        hvptest(D=D, G=G, z=z_mnist)

    print('=====Test2=====')
    D = DC_discriminator(channel_num=3, feature_num=64).to(device)
    G = DC_generator(z_dim=100, channel_num=3, feature_num=64).to(device)
    z_celeba = torch.randn((batch_size, 100, 1, 1), device=device)
    for i in range(10):
        hvptest(D=D, G=G, z=z_celeba)

</code></pre>
<p>Remember that you have to benchmark several iterations, cudnn.benchmark always slows down the first iteration.</p>

Hongkai_Zheng: <p>Oh, I see. thank you so much.</p>