How to enforce symmetry on generator output?

collinarnett: <p>I’m having a lot of issues implementing this paper, <a href="http://papers.nips.cc/paper/7978-generative-modeling-for-protein-structures.pdf" rel="noopener nofollow ugc">Generative Modeling for Protein Structures</a>. I’m using the <a href="https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html" rel="noopener nofollow ugc">DCGAN tutorial</a> as a reference since their architectures are similar.</p>
<h3><a name="context-1" class="anchor" href="#context-1"></a>Context:</h3>
<p>The dataloader in the tutorial is implemented as:</p>
<pre data-code-wrap="python"><code class="lang-python">dataset = dset.ImageFolder(...)
for i, data in enumerate(dataloader, 0):
    real_cpu = data[0].to(device)
</code></pre>
<p>and gives shapes:</p>
<pre><code class="lang-auto">print(data[0].shape)
print(data[1].shape)
torch.Size([128, 3, 64, 64])
torch.Size([128])
</code></pre>
<p>Whereas my dataset is implemented as:</p>
<pre><code class="lang-auto">with h5py.File('/home/collin/protein_maps/dataset.hdf5', 'r') as f:
    x = f['train_64'][:]
dataloader = torch.utils.data.DataLoader(x, ...)

for i, data in enumerate(dataloader, 0):
    # Unsqueezed dim one to convert [128, 64, 64] to [128, 1, 64, 64] to conform to D architecture 
    real_cpu = (data.unsqueeze(dim=1).type(torch.FloatTensor)).to(device)
</code></pre>
<p>This causes problems later on when I want to enforce symmetry later when passing from the generator to the discriminator as specified in the original paper:</p>
<blockquote>
<p>During training, we enforce that G(z) be positive by clamping output values above zero and symmetric by setting G(z) = (G(z)+G(z).T)/2 before passing the generated map to the discriminator.</p>
</blockquote>
<p>but I get a broadcasting error when trying to broadcast [128, 1, 64, 64] and [64, 64, 1, 128] which makes sense. The authors specify that this step be done in the generator architecture verbatim:</p>
<pre><code class="lang-auto">  Model architectures. Each layer is presented as:
   
  Layer(filters, kernel size, stride, padding)

  ------------------64 GAN------------------
  down-scale factor = 100
  
  --Generator--
  nz = 100
  ConvTranspose2d( 512, 4, 1, 0)
  BatchNorm2d(512)
  LeakyReLU(0.2),
  ConvTranspose2d(256, 4, 2, 1)
  BatchNorm2d(256)
  LeakyReLU(0.2)
  ConvTranspose2d(128, 4, 2, 1)
  BatchNorm2d(128)
  LeakyReLU(0.2)
  ConvTranspose2d(64, 4, 2, 1)
  BatchNorm2d(64)
  LeakyReLU(0.2)
  ConvTranspose2d(1, 4, 2, 1)
  Clamp(&gt;0)
  Enforce Symmetry
</code></pre>
<p>but I’m not sure how to do this in practice as my architecture looks like:</p>
<pre data-code-wrap="python"><code class="lang-python">self.main = nn.Sequential(
    nn.ConvTranspose2d(nz, 512, kernel_size=4, stride=1, padding=0),
    nn.BatchNorm2d(512),
    nn.LeakyReLU(0.2),
    nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),
    nn.BatchNorm2d(256),
    nn.LeakyReLU(0.2),
    nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
    nn.BatchNorm2d(128),
    nn.LeakyReLU(0.2),
    nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
    nn.BatchNorm2d(64),
    nn.ConvTranspose2d(64, 1, kernel_size=4, stride=2, padding=1),
)
</code></pre>
<h3><a name="question-2" class="anchor" href="#question-2"></a>Question</h3>
<p>How can I clamp and enfore symmetry in this part of my architecture? Is this possible or do I have to do this while training? If I do it in training what’s the best way to slice only the [64, 64] in [128, 1, 64, 64] and reinsert it after? Is there a function to enforce symmetry in pytorch in nn Sequential?</p>

ptrblck: <p>I’m not familiar with the paper, but wouldn’t this work:</p>
<pre><code class="lang-python">ret = (real_cpu + real_cpu.permute(0, 1, 3, 2)) / 2
</code></pre>

collinarnett: <p>Thank you for your reply. This worked but now it seems like my generator isn’t updating.</p>
<p>Here is my training loop:</p>
<pre><code class="lang-auto"># Training Loop
# Lists to keep track of progress
img_list = []
G_losses = []
D_losses = []
iters = 0

print("Starting Training Loop...")
# For each epoch
for epoch in range(num_epochs):
    # For each batch in the dataloader
    for i, data in enumerate(dataloader, 0):
        ############################
        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
        ###########################
        ## Train with all-real batch
        netD.zero_grad()
        # Format batch
        # Unsqueezed dim one to convert [128, 64, 64] to [128, 1, 64, 64] to conform to D architecture 
        real_cpu = (data.unsqueeze(dim=1).type(torch.FloatTensor)).to(device)
        b_size = real_cpu.size(0)
        label = torch.full((b_size,), real_label, device=device)
        # Forward pass real batch through D
        output = netD(real_cpu).view(-1)
        # Calculate loss on all-real batch
        errD_real = criterion(output, label)
        # Calculate gradients for D in backward pass
        errD_real.backward()
        D_x = output.mean().item()

        ## Train with all-fake batch
        # Generate batch of latent vectors
        noise = torch.randn(b_size, nz, 1, 1, device=device)
        # Generate fake image batch with G
        fake = netG(noise)
        label.fill_(fake_label)
        # Make Symmetric
        sym_fake = (fake.detach().clamp(min=0) + fake.detach().clamp(min=0).permute(0, 1, 3, 2)) / 2
        # Classify all fake batch with D
        output = netD(sym_fake).view(-1)
        # Calculate D's loss on the all-fake batch
        errD_fake = criterion(output, label)
        # Calculate the gradients for this batch
        errD_fake.backward()
        D_G_z1 = output.mean().item()
        # Add the gradients from the all-real and all-fake batches
        errD = errD_real + errD_fake
        # Update D
        optimizerD.step()
        #adjust_optim(optimizerD, iters)
        ############################
        # (2) Update G network: maximize log(D(G(z)))
        ###########################
        netG.zero_grad()
        label.fill_(real_label)  # fake labels are real for generator cost
        # Since we just updated D, perform another forward pass of all-fake batch through D
        output = netD(fake.detach()).view(-1)
        # Calculate G's loss based on this output
        errG = criterion(output, label)
        # Calculate gradients for G
        errG.backward()
        D_G_z2 = output.mean().item()
        # Update G
        optimizerG.step()
        adjust_optim(optimizerG, iters)

        # Output training stats
        if i % 50 == 0:
            print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
                  % (epoch, num_epochs, i, len(dataloader),
                     errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))

        # Save Losses for plotting later
        G_losses.append(errG.item())
        D_losses.append(errD.item())

        # Check how the generator is doing by saving G's output on fixed_noise
        if (iters % 500 == 0) or ((epoch == num_epochs-1) and (i == len(dataloader)-1)):
            with torch.no_grad():
                fake = netG(fixed_noise).detach().cpu()
            img_list.append(vutils.make_grid(fake, padding=2, normalize=True))

        iters += 1
</code></pre>
<p>I used your function on the output of the generator when feeding to the discriminator but my generator fails to learn after this.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/e/ded56bb1ecc8045daaac0e564788e4172979e508.png" alt="image" data-base62-sha1="vNhfvyHQL22B6m05i3WUlrMDJji" width="642" height="351"></p>