How do I train a network in which the output of one network depends on two other networks, and its loss depends on their losses as well?

8Gitbrix: <p>I’m trying to implement this paper <a href="https://arxiv.org/pdf/1911.08079.pdf" rel="noopener nofollow ugc">Two-Stream FCNs to Balance Content and Style for Style Transfer</a>, and here’s a picture of the architecture:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/0/9006a2c94b4c82c58894f7e08d1be0847e55e4c4.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/9006a2c94b4c82c58894f7e08d1be0847e55e4c4" title="Screen Shot 2020-11-26 at 2.20.00 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/0/9006a2c94b4c82c58894f7e08d1be0847e55e4c4_2_690x371.jpeg" alt="Screen Shot 2020-11-26 at 2.20.00 PM" data-base62-sha1="ky6YkrzZBy2QoZCI5gpqMeCkFSc" width="690" height="371" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/0/9006a2c94b4c82c58894f7e08d1be0847e55e4c4_2_690x371.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/9/0/9006a2c94b4c82c58894f7e08d1be0847e55e4c4_2_1035x556.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/0/9006a2c94b4c82c58894f7e08d1be0847e55e4c4_2_1380x742.jpeg 2x" data-dominant-color="E4E6E8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2020-11-26 at 2.20.00 PM</span><span class="informations">1722×928 279 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The outputs of the content and style subnets are fed into the generator subnet.</p>
<p>The content subnet’s loss depends on the generator subnet’s output and the input content image, and the style subnet’s loss depends on the generator subnet’s output and the input style image.</p>
<p>But the generator’s loss is a linear combination of the style and content loss.<br>
The paper mentions training the networks simultaneously and updating each of the weights independently.</p>
<p>How would I go about training these networks simultaneously?</p>
<p>I’ve defined the optimizers and criterions of the 3 subnets. From my outputs the networks only seems to be learning the style of the style image. Is this the wrong approach?</p>
<pre><code class="lang-auto">
for i in range(epochs):

  for xs, _ in train_dataloader:

    s_optimizer.zero_grad()
    c_optimizer.zero_grad()
    g_optimizer.zero_grad()

    s_image, p1, p2, p3 = s_model(style_image)
    c_image = c_model(xs.to(device), gamma, p1, p2, p3)
    generated_images = g_model(c_image.clone(), s_image.clone(), gamma)
    
    # use pretrained vgg as a feature extractor
    generated_outputs = VGG(generated_images)
    
    style_loss = s_criterion(generated_outputs, style_image)
    content_loss = c_criterion(generated_images, xs.to(device))
    generator_loss = g_criterion(style_loss, content_loss)

    generator_loss.backward(retain_graph = True)
    style_loss.backward(retain_graph = True)
    content_loss.backward()
   
    g_optimizer.step()
    c_optimizer.step()
    s_optimizer.step()
</code></pre>